diff --git a/docs/qwen38_adaptation.rst b/docs/qwen38_adaptation.rst new file mode 100644 index 0000000000..d4bc2d4f9b --- /dev/null +++ b/docs/qwen38_adaptation.rst @@ -0,0 +1,48 @@ +Qwen3.8-Flash-Next +================== + +Qwen3.8-Flash-Next checkpoints use the ``qwen4_exp`` model type. The model +implementation combines gated residuals, n-gram embeddings, QSA and GDN +attention, with dedicated multimodal, MTP and cache-state integration. +Reusable operations live in ``lightllm/common/gated_residual``, +``lightllm/common/ngram_embedding`` and ``lightllm/common/qsa``. + +CPU embedding offload is disabled by default. The optional QSA TopK selector +is enabled by the model only for its qualified H100/Torch/Triton geometry; +other configurations use Torch. Shared-prefix admission is limited to the +model's single-node, DP1 normal serving mode with one-page state blocks. + +Validation +---------- + +Historical H100 measurements and their configuration are recorded in +``qwen38_h100_performance.rst``. They cover TP8 BF16, FP32 SSM states, +fixed MTP2 and one-page linear-attention state blocks. These results do not +qualify PD/cross-DP serving or other cache block configurations. + +The four-page-block pressure-recovery configuration has unresolved text +differences. Low-load latency still trails SGLang in some metrics. Archived +HC and sampling experiments are not part of the runtime implementation. + +Rebase to upstream, 2026-11-03 +------------------------------ + +The implementation now uses AttentionStateManager. QSA keys and positions are +owned by the model's GQA State and participate in token copies and CPU/PD/DP +transfer layouts. Lexical histories use a separate bounded State for reset, +prefix snapshots, CPU/PD/DP transfers and accepted MTP candidate commits. +The removed request/cache manager classes are not restored. + +Model and draft registration use the upstream lazy registry. HTTP slot waiting +preserves atomic group allocation and current PD timeouts. The GDN MTP path uses +the upstream recurrent kernel and its explicit run configuration. + +The historical H100 results above predate this interface migration. They do not +validate the rebased runtime. GPU generation, complete PD transfer and performance +must be rechecked before deploying a new image of this revision. + +CPU validation after migration: 1022 checks passed and 930 GPU cases were +skipped. The DeepSeek V4 integration module could not be collected because +FlashMLA is absent. Two upstream State identity cases require CUDA-backed +pinned memory; both failures reproduce on unmodified upstream/main with CUDA +hidden. No GPU serving, PD end-to-end or performance result is claimed here. diff --git a/docs/qwen38_h100_performance.rst b/docs/qwen38_h100_performance.rst new file mode 100644 index 0000000000..bb9bd4be39 --- /dev/null +++ b/docs/qwen38_h100_performance.rst @@ -0,0 +1,76 @@ +Qwen3.8-Flash-Next H100 性能与精度 +================================= + +当前结论：4× 的两轮吞吐、TTFT P95、完整 HTTP P95 和末次可见文本 P95 均优于冻结基线。 +低负载下仍有延迟差距，尚未完成“超过 vLLM 和 SGLang”的整体目标。 + +测试口径 +-------- + +* 8 × NVIDIA H100 80GB HBM3，TP8，BF16，SSM 状态 FP32；MTP 固定 2。 +* 模型版本 de4b8e4d43b917e7706784d8bb445c9af86a3540；4096 上下文、最多 8 个并发请求、512-token prefill chunk、16384 KV 容量。 +* 当前已验证的线性注意力状态页为 512 tokens，每块一页。 +* AIPerf 0.9 统一回放公开 Mooncake 轨迹中的 256 个请求，倍率为 1×、2×、4×。 +* 公开轨迹提供匿名长度、时间戳和 hash；请求正文由 AIPerf 合成，并非原始私有正文。 +* 每个 LightLLM 样本使用全新服务；vLLM、SGLang 只使用此前冻结的原始结果，未重新推理。 +* 原生 Torch 2.9.1+cu128 / Triton 3.5.1；vLLM 0.31；SGLang 0.5.22。 +* 1×/2× 吞吐受请求到达速率限制，不能当作最大处理容量；保留所有重复样本。 + +已完成的性能样本 +---------------- + +.. csv-table:: 延迟单位 ms；吞吐单位成功输出 tok/s + :header: "倍率", "框架", "样本", "吞吐 ↑", "TTFT P95 ↓", "完整 HTTP P95 ↓", "末次文本 P95 ↓", "空文本错误/256" + + "1×", "LightLLM", "1", "276.14", "516.09", "3832.07", "3615.36", "2" + "1×", "LightLLM", "2", "276.16", "515.16", "3818.23", "3622.51", "2" + "1×", "SGLang", "冻结", "276.99", "610.40", "3864.63", "3329.87", "2" + "1×", "vLLM", "冻结", "277.72", "804.93", "4255.93", "3548.29", "1" + "2×", "LightLLM", "1", "546.33", "475.34", "4466.54", "4187.74", "2" + "2×", "SGLang", "冻结", "545.80", "600.51", "5302.52", "4178.47", "2" + "2×", "vLLM", "冻结", "513.20", "2933.67", "8695.77", "7227.73", "2" + "4×", "LightLLM", "1", "887.26", "376.17", "5794.25", "5188.47", "2" + "4×", "LightLLM", "2", "886.14", "390.13", "5770.75", "5151.58", "2" + "4×", "SGLang", "冻结", "835.28", "864.93", "5924.62", "5204.95", "2" + "4×", "vLLM", "冻结", "523.95", "3367.76", "9085.75", "7930.38", "1" + +完整 HTTP 统计包含全部 256 个请求；TTFT 和末次文本仅包含有可见文本的请求。 +空文本仍计为错误，其输出 token 不计入成功吞吐。AIPerf 导出的 request_latency +在这些记录中对应末次可见文本，不能替代完整 HTTP/DONE 边界。 + +精度与兼容性 +------------ + +* 数学题正确 1272/1319，准确率 96.44%。 +* 全部 1319 个数学样例的生成 token 序列与原实现一致；分别保留 19 例截断、19 例无有效最终答案。 +* 480 个冷缓存请求和 48 个调度、缓存、EOS 请求核对通过；2788 个训练模型中间张量完全一致。 +* H100 共享组件/默认行为回归 134 项通过；本地 CPU 兼容性检查 124 项通过，零失败、错误、跳过。 +* 已完成的 5 轮回放共 1280 个输出文本与原实现一致，空文本错误单独保留。 +* TopK 为独立组件，共享选项默认关闭，Qwen4 在工厂启用，select 接口保持兼容。 + +剩余工作与限制 +-------------- + +* 1× 首轮末次文本 P95 为 3615.36 ms，高于 SGLang 3329.87 ms；中位请求延迟也仍落后。 +* 2× 末次文本 P95 为 4187.74 ms，比 SGLang 4178.47 ms 高约 9.28 ms。 +* 历史四页块压力恢复的文本差异仍未解决，不能将当前一页块精度结论推广到该配置。 +* 五轮回放及低负载剖析已完成；按用户要求收尾，优化目标已暂停，未标记为全面达成。 + +原始结果与证据 +-------------- + +* 五轮独立汇总：``artifacts/h100/261019/audits/topk-short-five-sample-summary.json``。 +* 性能 CSV：``artifacts/h100/261019/audits/topk-short-performance-current.csv``。 +* 完整精度：``artifacts/h100/261019/audits/topk-short-full-math-token-parity.json``。 +* 源码与本地验证：``artifacts/h100/261019/audits/topk-short-local-adoption.json``。 +* 重复回放状态：``artifacts/h100/261019/tuning/topk-short-canonical-candidate/progress.json``。 + +收尾状态 +-------- + +* 最终保留完整精度和五轮回放验证通过的 TopK 版本；H100 隔离源码已恢复到同一版本。 +* HC 小批量上投影候选已归档：30 项 CPU、203 项 GPU 检查通过，2788 个模型中间张量一致。 + 完整模型轨迹实际覆盖 8 行 BD8；其余选择形状只有组件检查。HC 尚未完成服务精度与公开回放验证，未纳入最终实现或上述性能结果。 +* 首次 HC 模型轨迹因测试插桩作用域断言失败而回滚；修复插桩并验证后重跑通过，失败记录一并保留。 +* H100 恢复后核对 949 项源码和 42 项测试/基准文件哈希；GPU 无遗留计算进程。未启动新的推理测试。 +* 收尾记录：``artifacts/h100/261019/audits/user-wrapup-qualified-topk-restored.json``。 diff --git a/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py index 3fe6f2317e..82a07fc486 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py @@ -13,8 +13,24 @@ def create_gdn_state( request_capacity: int, mtp_step: int = 0, device: str | torch.device = "cuda", + ssm_state_mode: str = "native", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, ) -> GDNState: - """创建 GDN 请求运行态，大小页池由后续容量分配流程初始化。""" + """创建 GDN 请求运行态。replay 模式使用独立的 ReplayGDNState。""" + if ssm_state_mode == "replay": + from .replay_gdn_state import ReplayGDNState + + return ReplayGDNState( + config=config, + request_capacity=request_capacity, + mtp_step=mtp_step, + device=device, + ssm_replay_config=ssm_replay_config, + gate_lower_bound=gate_lower_bound, + ) + if ssm_state_mode != "native": + raise ValueError(f"unknown GDN state mode: {ssm_state_mode}") return GDNState( config=config, request_capacity=request_capacity, 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 new file mode 100644 index 0000000000..d927891ada --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -0,0 +1,835 @@ +"""GDN replay 的 bounded attention 状态。 + +Conv、canonical SSM、replay history、页布局和已接受快照都在本类。 +一层 decode 重放在 replay_gdn_decode。 +""" + +from __future__ import annotations + +from functools import lru_cache +from math import prod +from typing import NamedTuple + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( + _advance_cursor, + _fold_prefix, +) +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import _reset_folded_cursors +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.utils.log_utils import init_logger + +from ...types import AxisMapping, PageKind, StateTransferLayout, TensorField, TransferRegion +from ...utils import deepcopy_result, dtype_name +from ..base_state import BoundedAttentionState +from ..conv_ssm_types import ConvSSMStateConfig +from ..copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from .gdn_pages_pool import GDNPagesPool + +logger = init_logger(__name__) + + +class _AcceptedTensors(NamedTuple): + conv: torch.Tensor + ssm: torch.Tensor + + +class ReplayGDNState(BoundedAttentionState): + """GDN replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", + ssm_state_mode: str = "replay", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, + ) -> None: + """创建单槽 canonical SSM，并挂上只属于当前请求的 replay history。""" + if ssm_state_mode != "replay": + raise ValueError(f"ReplayGDNState only implements replay mode, got {ssm_state_mode}") + assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" + self.config = config + self.state_name = None + self.request_capacity = request_capacity + self.mtp_step = mtp_step + self.gate_lower_bound = gate_lower_bound + self._algorithm = "gdn" + self._logical_to_storage_layer = {layer_id: index for index, layer_id in enumerate(config.logical_layer_ids)} + layer_count = len(config.logical_layer_ids) + self.conv_state = torch.zeros( + ( + layer_count, + request_capacity, + config.conv_channels, + config.conv_width + mtp_step, + ), + dtype=config.conv_dtype, + device=device, + ) + self.ssm_state = torch.zeros( + ( + layer_count, + request_capacity, + config.num_v_heads, + config.key_dim, + config.value_dim, + ), + dtype=config.ssm_dtype, + device=device, + ) + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.big_pages_pool = None + self.small_pages_pool = None + replay_config = dict(ssm_replay_config or {}) + capacity = int(replay_config.get("cache_len", 16)) + compute_mode = replay_config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") + force_compress = bool(replay_config.get("force_compress", False)) + _bind_history( + self, + self.ssm_state, + capacity, + mtp_step + 1, + config.conv_dtype, + config.num_k_heads, + gate_lower_bound, + compute_mode, + None, + force_compress, + ) + + def get_state_type(self) -> str: + """返回算法名和 Conv/SSM 存储精度。""" + return ( + f"{self._algorithm}_conv_{dtype_name(dtype=self.config.conv_dtype)}" + f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" + ) + + def initialize_page_pools( + self, + big_page_capacity: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" + assert self.big_pages_pool is None and self.small_pages_pool is None, "page pools already initialized" + assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" + self.big_pages_pool = ( + GDNPagesPool( + size=big_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if big_page_capacity + else None + ) + self.small_pages_pool = ( + GDNPagesPool( + size=small_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if small_page_capacity + else None + ) + + # ======================================================================== + # 二、计算视图、运行态重置、已接受快照与 MTP 提交 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射。""" + return logical_layer_id in self._logical_to_storage_layer + + def storage_layer(self, logical_layer_id: int) -> int: + """返回逻辑层在 history 和 canonical state 里的存储下标。""" + + return self._logical_to_storage_layer[logical_layer_id] + + def get_layer_view(self, logical_layer_id: int): + index = self.storage_layer(logical_layer_id) + return self.conv_state[index], self.ssm_state[index] + + def reset_bounded_att_state(self, req_indices: list[int]) -> None: + """清零请求的 Conv、单槽 SSM、history 和 MTP 游标。""" + assert isinstance(req_indices, list), "request indices must be a list" + for req_idx in req_indices: + self.conv_state[:, req_idx].zero_() + self.ssm_state[:, req_idx].zero_() + self.clear_history(req_idx) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + def save_bounded_att_runtime_to_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind, + page_slot_ids: list[int], + ) -> None: + """把已接受 Conv 窗口和物化 SSM 写入页，不改 canonical state。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + copy_bounded_att_tensors( + copies=[ + (self._accepted_conv(req_idx), pool.conv_state[slot_id]) + for req_idx, slot_id in zip(req_indices, page_slot_ids) + ] + ) + self.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) + + def restore_bounded_att_runtime_from_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind, + page_slot_ids: list[int], + ) -> None: + """从页恢复基准 Conv/SSM，并清空目标请求的 history。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + width = self.config.conv_width + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + self.clear_history(req_idx) + conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, :width])) + ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + if self.mtp_state_index is not None: + for req_idx in req_indices: + self.mtp_state_index[req_idx].fill_(0) + + def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: + """提交已接受长度，再按容量或 force_compress 决定是否整段压缩。""" + if self.mtp_state_index is None: + return + self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) + self.maybe_compress_prefix(req_indices) + + def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: + """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" + return _AcceptedTensors(self._accepted_conv(req_idx), self.accepted_ssm(req_idx)) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回已接受 Conv/SSM 快照的字段区间。不含 history 或 MTP 候选区。""" + return self._transfer_regions() + + def write_state_pages_to_cpu( + self, + page_kind: PageKind, + page_slot_ids: list[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """将指定页池槽位的快照写入 CPU 页。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self._pool(page_kind) + assert isinstance(page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(page_slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" + if cpu_page_ready_flags is not None: + assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" + page_fields, page_indices = [], [] + for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): + # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 + if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: + continue + page_fields.append(self._tensor_fields(pool.conv_state[slot], pool.ssm_state[slot])) + page_indices.append(page) + for copies in bounded_att_buffer_copies( + page_fields=page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=page_indices, + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + def read_state_pages_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + big_page_slot_ids: list[int], + layout: StateTransferLayout, + ) -> None: + """读入指定大页槽位。请求运行态由 restore_bounded_att_runtime_from_pages 另行恢复。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self.big_pages_pool + assert isinstance(big_page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" + # 本地目标保留本 TP rank 的分片形状。layout.regions 描述 CPU 全局字段。 + pool_page_fields = [ + self._tensor_fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids + ] + for copies in bounded_att_buffer_copies( + page_fields=pool_page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=cpu_page_indices, + write=False, + ): + copy_bounded_att_tensors(copies=copies) + + # ======================================================================== + # 四、DP 迁移：大页快照与基准运行态 + # ======================================================================== + + def copy_big_pages_from_dp(self, src_page_ids: list[int], dst_page_ids: list[int], remote_dp_bounded_state) -> None: + """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" + assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" + assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + src_pool = remote_dp_bounded_state.big_pages_pool + dst_pool = self.big_pages_pool + assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" + # 大页是基准快照，与请求槽位无关。页池须在本进程注册，当前流访问完才能归还页面。 + conv_copies, ssm_copies = [], [] + for src_page, dst_page in zip(src_page_ids, dst_page_ids): + conv_copies.append((src_pool.conv_state[src_page], dst_pool.conv_state[dst_page])) + ssm_copies.append((src_pool.ssm_state[src_page], dst_pool.ssm_state[dst_page])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + + def copy_runtime_state_from_dp( + self, + src_req_indices: list[int], + dst_req_indices: list[int], + remote_dp_bounded_state: "ReplayGDNState", + ) -> None: + """复制已接受运行态到目标单槽，并清空目标 history。""" + assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" + assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + width = self.config.conv_width + conv_copies = [] + for src_req, dst_req in zip(src_req_indices, dst_req_indices): + self.clear_history(dst_req) + conv_copies.append( + ( + remote_dp_bounded_state._accepted_conv(src_req), + self.conv_state[:, dst_req, :, :width], + ) + ) + copy_bounded_att_tensors(copies=conv_copies) + remote_dp_bounded_state.export_ssm(src_req_indices, [self.ssm_state[:, dst_req] for dst_req in dst_req_indices]) + if self.mtp_state_index is not None: + for dst_req in dst_req_indices: + self.mtp_state_index[dst_req].fill_(0) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" + return self._transfer_regions() + + def write_to_pd_page( + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """写出已接受 Conv 窗口和物化 SSM，不改写源 history。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + accepted = self.get_accepted_state(req_idx) + fields = self._tensor_fields(accepted.conv, accepted.ssm) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """读入基准 Conv/SSM，并清空该请求的 history。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + self.clear_history(req_idx) + fields = self._tensor_fields( + self.conv_state[:, req_idx, :, : self.config.conv_width], + self.ssm_state[:, req_idx], + ) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=False, + ): + copy_bounded_att_tensors(copies=copies) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + # ======================================================================== + # 六、Replay history + # ======================================================================== + + def finish_prefill(self, logical_layer_id: int, reqs: torch.Tensor) -> None: + """最后一层把已接受前缀整段压进 SSM。""" + + layers = () if self.config is None else self.config.logical_layer_ids + if layers and logical_layer_id == layers[-1]: + self.compress_prefix(reqs) + + def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None: + """最后一层记下这 1 个 token，再按容量决定是否整段压缩。""" + + layers = () if self.config is None else self.config.logical_layer_ids + if not layers or logical_layer_id != layers[-1]: + return + self.accept_single_token(reqs) + self.maybe_compress_prefix(reqs) + + def clear_history(self, req: int | torch.Tensor) -> None: + """清空复用或恢复请求的临时 history，保留 canonical SSM。""" + + if type(req) is int: + # 标量赋值会把 CPU tensor 复制到 CUDA slot。 + self.cursors[req].zero_() + else: + self.cursors[req] = 0 + + def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: + """返回本步写入前的前缀长度。decode 不在这里修改 SSM。""" + + del cu_seqlens + return self.cursors.index_select(0, reqs.long()) + + def accept_updates( + self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None, width: int | None = None + ) -> None: + """更新 cursor。width 缺省时按 verify_width 取已接受长度，传 1 时固定加一个 token。""" + + _advance_cursor[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + accepted_state_indices, + reqs.numel(), + self.hold_req_id, + self.verify_width if width is None else width, + 256, + ) + + def accept_single_token(self, reqs: torch.Tensor) -> None: + """把一个已接受 token 记入前缀长度。""" + + if reqs.numel() == 0: + return + # width=1 不读 index。仍传 reqs，kernel 不能收到空指针。 + self.accept_updates(reqs, reqs, width=1) + + def merge_accepted_updates(self, reqs: torch.Tensor) -> None: + """把 cursor 覆盖的已接受前缀整段折回 canonical SSM，并清空 cursor。""" + + if reqs.numel() == 0: + return + self._materialize_accepted(reqs, self.ssm_state, snapshot=False) + self.cursors.index_fill_(0, reqs.long(), 0) + + def compress_prefix(self, reqs: torch.Tensor) -> None: + """把 cursor 覆盖的已接受前缀整段压进 SSM。长度为 0 时 kernel 直接返回。""" + + if reqs.numel() == 0: + return + if self.ssm_state.is_cuda: + self.merge_accepted_updates(reqs) + else: + self.clear_history(reqs) + + def maybe_compress_prefix(self, reqs: torch.Tensor) -> None: + """长度加上一次最大 verify_width 放不下，或 force_compress 打开时，整段压缩。""" + + if reqs.numel() == 0 or not self.ssm_state.is_cuda: + return + if self.force_compress: + self.compress_prefix(reqs) + return + self._materialize_accepted(reqs, self.ssm_state, snapshot=False, when_full=True) + + def accepted_ssm(self, req_idx: int) -> torch.Tensor: + """返回不修改 canonical SSM 和 history 的单请求快照。""" + + layers, _, hv, k, v = self.ssm_state.shape + output = torch.empty((layers, hv, k, v), dtype=self.ssm_state.dtype, device=self.ssm_state.device) + reqs = torch.tensor([req_idx], dtype=torch.int32, device=self.ssm_state.device) + self._materialize_accepted(reqs, output, snapshot=True) + return output + + def export_accepted_state( + self, + req_indices: list[int], + targets: list[torch.Tensor], + flush: bool = False, + ) -> None: + """将多个 accepted 快照直接写入调用方目标页。""" + + if not req_indices: + return + shape = (self.ssm_state.shape[0], *self.ssm_state.shape[2:]) + strides = targets[0].stride() + assert len(req_indices) == len(targets) + assert all(t.shape == shape and t.stride() == strides and t.dtype == self.ssm_state.dtype for t in targets) + device = self.ssm_state.device + self._materialize_accepted( + torch.tensor(req_indices, dtype=torch.int32, pin_memory=True).to(device, non_blocking=True), + torch.tensor([t.data_ptr() for t in targets], dtype=torch.int64, pin_memory=True).to( + device, non_blocking=True + ), + snapshot=not flush, + flush_export=flush, + output_strides=strides, + ) + + def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: + """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" + + if self.verify_width > 1: + commit_ssm_prefix(indices, reqs, lengths, self.cursors) + else: + commit_ssm_prefix(indices, reqs, lengths) + self.accept_updates(reqs, indices) + + def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: bool = False) -> None: + """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先物化再拷贝。""" + + if not req_indices: + return + target = targets[0] + if target.is_cuda and target.device == self.ssm_state.device: + self.export_accepted_state(req_indices, targets, flush=flush) + return + for req, dst in zip(req_indices, targets): + snap = self.accepted_ssm(req) + dst.copy_(snap) + if flush: + self.ssm_state[:, req].copy_(snap) + self.clear_history(req) + + _run_config: dict[str, int] | None = None + + @property + def run_config(self) -> dict[str, int] | None: + """返回物化和 autotune 使用的 Triton 配置，不是 SSM 布局。""" + + return self._run_config + + @run_config.setter + def run_config(self, config: dict | None) -> None: + """保存一份与调用方脱开的 kernel 配置。""" + + if config is None: + self._run_config = None + return + self._run_config = { + "BV": int(config["BV"]), + "num_warps": int(config["num_warps"]), + "num_stages": int(config.get("num_stages", 3)), + } + + # ======================================================================== + # 七、私有辅助实现 + # ======================================================================== + + def _transfer_regions(self) -> tuple[TransferRegion, ...]: + """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" + # 只传有效历史窗口，不包含 MTP 候选区和游标。 + config = self.config + storage_layer_count = len(config.logical_layer_ids) + conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim + conv_shape = (storage_layer_count, conv_channels, config.conv_width) + conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, Conv shape=%s", + self.state_name, + config.tp_rank, + config.tp_world_size, + conv_shape, + ) + + # Conv 按 Q/K/V 拼接。三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 + local_channel_start = 0 + global_channel_start = 0 + for name, local_head_count, global_head_count, head_dim in ( + ("Q", config.num_k_heads, config.global_k_heads, config.key_dim), + ("K", config.num_k_heads, config.global_k_heads, config.key_dim), + ("V", config.num_v_heads, config.global_v_heads, config.value_dim), + ): + # 每组连续 rank 保存相同 heads，仅组内首 rank 写出。 + replica_count = config.tp_world_size * local_head_count // global_head_count + head_start = config.tp_rank // replica_count * local_head_count + write_enabled = config.tp_rank % replica_count == 0 + local_channel_count = local_head_count * head_dim + target_channel_start = global_channel_start + head_start * head_dim + regions.append( + TransferRegion( + field="conv", + byte_start=0, + byte_end=conv_byte_end, + dtype=config.conv_dtype, + axes=("layer", "channel", "window"), + shape=conv_shape, + axis_mappings=( + AxisMapping( + axis="channel", + local_start=local_channel_start, + external_start=target_channel_start, + length=local_channel_count, + ), + AxisMapping(axis="window", local_start=0, external_start=0, length=config.conv_width), + ), + write_enabled=write_enabled, + ) + ) + logger.debug( + "%s Conv %s: local_channel_start=%d, global_channel_start=%d, " + "channels=%d, replicas=%d, write_enabled=%s", + self.state_name, + name, + local_channel_start, + target_channel_start, + local_channel_count, + replica_count, + write_enabled, + ) + local_channel_start += local_channel_count + global_channel_start += global_head_count * head_dim + + # SSM 沿 V head 切分，接在 Conv 之后，并按 SSM dtype 对齐字节起点。 + replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads + head_start = config.tp_rank // replica_count * config.num_v_heads + write_enabled = config.tp_rank % replica_count == 0 + ssm_shape = (storage_layer_count, config.global_v_heads, config.key_dim, config.value_dim) + element_bytes = config.ssm_dtype.itemsize + ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes + ssm_region = TransferRegion( + field="ssm", + byte_start=ssm_byte_start, + byte_end=ssm_byte_start + prod(ssm_shape) * element_bytes, + dtype=config.ssm_dtype, + axes=("layer", "head", "key_dim", "value_dim"), + shape=ssm_shape, + axis_mappings=( + AxisMapping(axis="head", local_start=0, external_start=head_start, length=config.num_v_heads), + AxisMapping(axis="key_dim", local_start=0, external_start=0, length=config.key_dim), + AxisMapping(axis="value_dim", local_start=0, external_start=0, length=config.value_dim), + ), + write_enabled=write_enabled, + ) + regions.append(ssm_region) + logger.debug( + "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", + self.state_name, + head_start, + config.num_v_heads, + write_enabled, + ssm_region.byte_start, + ssm_region.byte_end, + ) + return tuple(regions) + + def _pool(self, page_kind: PageKind): + """根据大小页类型取得已启用的页池，不访问 allocator。""" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" + return pool + + @staticmethod + def _tensor_fields(conv: torch.Tensor, ssm: torch.Tensor): + """描述一页里的 Conv 和 SSM 字段。""" + return { + "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), + "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), + } + + def _accepted_conv(self, req_idx: int) -> torch.Tensor: + """取出该请求已接受的 Conv 窗口，不包含 MTP 候选区。""" + width = self.config.conv_width + conv = self.conv_state[:, req_idx] + if self.mtp_state_index is None: + return conv[:, :, :width] + offset = self.mtp_state_index[req_idx] + return conv.index_select(-1, torch.arange(width, device=conv.device) + offset) + + def _materialize_accepted( + self, + reqs: torch.Tensor, + output: torch.Tensor, + snapshot: bool, + output_strides=None, + flush_export: bool = False, + when_full: bool = False, + ) -> None: + """将请求的 accepted prefix 写入 canonical SSM 或调用方提供的输出页。""" + + assert not (snapshot and flush_export), "flush export must update the canonical state" + layers, slots, hv, k, v = self.ssm_state.shape + config = self.run_config or {"BV": 32, "num_warps": 4, "num_stages": 3} + bv = config["BV"] + _fold_prefix[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( + self.ssm_state, + output, + self.raw_keys, + self.raw_values, + self.gates, + self.betas, + self.cursors, + reqs, + self.ssm_state.stride(), + slots, + hv, + self.num_key_heads, + k, + v, + self.capacity, + self.hold_req_id, + triton.next_power_of_2(k), + bv, + snapshot, + False, + FLUSH_EXPORT=flush_export, + OUTPUT_STRIDES=output_strides, + WHEN_FULL=when_full, + NEXT=self.verify_width, + num_warps=config["num_warps"], + num_stages=config["num_stages"], + ) + if flush_export or when_full: + _reset_folded_cursors[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + reqs.numel(), + self.hold_req_id, + 0 if flush_export else self.capacity, + self.verify_width, + 256, + ) + + +def _history_owner( + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype = torch.bfloat16, + num_key_heads: int | None = None, + gate_lower_bound: float = -5.0, + compute_mode: str = "inline", + run_config: dict[str, int] | None = None, +) -> ReplayGDNState: + """把 history 挂到调用方已有的 SSM 上，不再分配第二份状态。供保持原 stride 的 kernel 测试使用。""" + + owner = ReplayGDNState.__new__(ReplayGDNState) + owner.config = None + owner.state_name = None + owner.request_capacity = ssm.shape[1] + owner.mtp_step = verify_width - 1 + owner._algorithm = "gdn" + owner._logical_to_storage_layer = {} + owner.conv_state = None + owner.mtp_state_index = None + owner.big_pages_pool = None + owner.small_pages_pool = None + _bind_history( + owner, + ssm, + capacity, + verify_width, + activation_dtype, + num_key_heads, + gate_lower_bound, + compute_mode, + run_config, + ) + return owner + + +def _bind_history( + owner, + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype, + num_key_heads: int | None, + gate_lower_bound: float, + compute_mode: str, + run_config: dict[str, int] | None, + force_compress: bool = False, +) -> None: + """在 owner 上创建 GDN replay 的连续前缀 history 和 cursor。""" + + assert ssm.dtype in (torch.float32, torch.bfloat16), "GDN replay requires FP32 or BF16 SSM state" + assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 + assert compute_mode in ("inline", "precompute") + owner.ssm_state = ssm + owner.kda = False + owner.gate_lower_bound = gate_lower_bound + owner.capacity = capacity + owner.verify_width = verify_width + owner.force_compress = bool(force_compress) + owner.compute_mode = compute_mode + owner.run_config = run_config + owner._config_is_fixed = run_config is not None + owner.batch_configs = {} + layers, slots, hv, k, v = ssm.shape + owner.num_key_heads = hv if num_key_heads is None else num_key_heads + assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 + owner.hold_req_id = slots - 1 + owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) + # 归一化 k 可用 fp16；delta 保留 fp32，避免有限 BF16 输入的修正量溢出 fp16。 + # 原始 k/v 保持激活精度。g 和 beta 保持 fp32。压缩在 fp32 里从原始值折回。 + reconstruction_dtype = torch.float16 if activation_dtype == torch.bfloat16 else torch.float32 + key_shape = (layers, slots, owner.num_key_heads, capacity, k) + owner.keys = torch.empty(key_shape, dtype=reconstruction_dtype, device=ssm.device) + owner.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=torch.float32, device=ssm.device) + owner.gates = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) + owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) + owner.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=activation_dtype, device=ssm.device) + owner.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) diff --git a/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py index 49c9ad0ce5..1742441389 100644 --- a/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py @@ -13,8 +13,24 @@ def create_kda_state( request_capacity: int, mtp_step: int = 0, device: str | torch.device = "cuda", + ssm_state_mode: str = "native", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, ) -> KDAState: - """创建 KDA 请求运行态，大小页池由后续容量分配流程初始化。""" + """创建 KDA 请求运行态。replay 模式使用独立的 ReplayKDAState。""" + if ssm_state_mode == "replay": + from .replay_kda_state import ReplayKDAState + + return ReplayKDAState( + config=config, + request_capacity=request_capacity, + mtp_step=mtp_step, + device=device, + ssm_replay_config=ssm_replay_config, + gate_lower_bound=gate_lower_bound, + ) + if ssm_state_mode != "native": + raise ValueError(f"unknown KDA state mode: {ssm_state_mode}") return KDAState( config=config, request_capacity=request_capacity, 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 new file mode 100644 index 0000000000..c2507e03e9 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -0,0 +1,834 @@ +"""KDA replay 的 bounded attention 状态。 + +Conv、canonical SSM、replay history、页布局和已接受快照都在本类。 +一层 decode 重放在 replay_kda_decode。 +""" + +from __future__ import annotations + +from functools import lru_cache +from math import prod +from typing import NamedTuple + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( + _advance_cursor, + _fold_prefix, +) +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import _reset_folded_cursors +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.utils.log_utils import init_logger + +from ...types import AxisMapping, PageKind, StateTransferLayout, TensorField, TransferRegion +from ...utils import deepcopy_result, dtype_name +from ..base_state import BoundedAttentionState +from ..conv_ssm_types import ConvSSMStateConfig +from ..copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from .kda_pages_pool import KDAPagesPool + +logger = init_logger(__name__) + + +class _AcceptedTensors(NamedTuple): + conv: torch.Tensor + ssm: torch.Tensor + + +class ReplayKDAState(BoundedAttentionState): + """KDA replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", + ssm_state_mode: str = "replay", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, + ) -> None: + """创建单槽 canonical SSM，并挂上只属于当前请求的 replay history。""" + if ssm_state_mode != "replay": + raise ValueError(f"ReplayKDAState only implements replay mode, got {ssm_state_mode}") + assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" + self.config = config + self.state_name = None + self.request_capacity = request_capacity + self.mtp_step = mtp_step + self.gate_lower_bound = gate_lower_bound + self._algorithm = "kda" + self._logical_to_storage_layer = {layer_id: index for index, layer_id in enumerate(config.logical_layer_ids)} + layer_count = len(config.logical_layer_ids) + self.conv_state = torch.zeros( + ( + layer_count, + request_capacity, + config.conv_channels, + config.conv_width + mtp_step, + ), + dtype=config.conv_dtype, + device=device, + ) + self.ssm_state = torch.zeros( + ( + layer_count, + request_capacity, + config.num_v_heads, + config.key_dim, + config.value_dim, + ), + dtype=config.ssm_dtype, + device=device, + ) + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.big_pages_pool = None + self.small_pages_pool = None + replay_config = dict(ssm_replay_config or {}) + capacity = int(replay_config.get("cache_len", 16)) + compute_mode = replay_config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") + force_compress = bool(replay_config.get("force_compress", False)) + _bind_history( + self, + self.ssm_state, + capacity, + mtp_step + 1, + config.conv_dtype, + config.num_k_heads, + gate_lower_bound, + compute_mode, + None, + force_compress, + ) + + def get_state_type(self) -> str: + """返回算法名和 Conv/SSM 存储精度。""" + return ( + f"{self._algorithm}_conv_{dtype_name(dtype=self.config.conv_dtype)}" + f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" + ) + + def initialize_page_pools( + self, + big_page_capacity: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" + assert self.big_pages_pool is None and self.small_pages_pool is None, "page pools already initialized" + assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" + self.big_pages_pool = ( + KDAPagesPool( + size=big_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if big_page_capacity + else None + ) + self.small_pages_pool = ( + KDAPagesPool( + size=small_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if small_page_capacity + else None + ) + + # ======================================================================== + # 二、计算视图、运行态重置、已接受快照与 MTP 提交 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射。""" + return logical_layer_id in self._logical_to_storage_layer + + def storage_layer(self, logical_layer_id: int) -> int: + """返回逻辑层在 history 和 canonical state 里的存储下标。""" + + return self._logical_to_storage_layer[logical_layer_id] + + def get_layer_view(self, logical_layer_id: int): + index = self.storage_layer(logical_layer_id) + return self.conv_state[index], self.ssm_state[index] + + def reset_bounded_att_state(self, req_indices: list[int]) -> None: + """清零请求的 Conv、单槽 SSM、history 和 MTP 游标。""" + assert isinstance(req_indices, list), "request indices must be a list" + for req_idx in req_indices: + self.conv_state[:, req_idx].zero_() + self.ssm_state[:, req_idx].zero_() + self.clear_history(req_idx) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + def save_bounded_att_runtime_to_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind, + page_slot_ids: list[int], + ) -> None: + """把已接受 Conv 窗口和物化 SSM 写入页，不改 canonical state。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + copy_bounded_att_tensors( + copies=[ + (self._accepted_conv(req_idx), pool.conv_state[slot_id]) + for req_idx, slot_id in zip(req_indices, page_slot_ids) + ] + ) + self.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) + + def restore_bounded_att_runtime_from_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind, + page_slot_ids: list[int], + ) -> None: + """从页恢复基准 Conv/SSM，并清空目标请求的 history。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + width = self.config.conv_width + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + self.clear_history(req_idx) + conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, :width])) + ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + if self.mtp_state_index is not None: + for req_idx in req_indices: + self.mtp_state_index[req_idx].fill_(0) + + def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: + """提交已接受长度，再按容量或 force_compress 决定是否整段压缩。""" + if self.mtp_state_index is None: + return + self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) + self.maybe_compress_prefix(req_indices) + + def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: + """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" + return _AcceptedTensors(self._accepted_conv(req_idx), self.accepted_ssm(req_idx)) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回已接受 Conv/SSM 快照的字段区间。不含 history 或 MTP 候选区。""" + return self._transfer_regions() + + def write_state_pages_to_cpu( + self, + page_kind: PageKind, + page_slot_ids: list[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """将指定页池槽位的快照写入 CPU 页。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self._pool(page_kind) + assert isinstance(page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(page_slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" + if cpu_page_ready_flags is not None: + assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" + page_fields, page_indices = [], [] + for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): + # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 + if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: + continue + page_fields.append(self._tensor_fields(pool.conv_state[slot], pool.ssm_state[slot])) + page_indices.append(page) + for copies in bounded_att_buffer_copies( + page_fields=page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=page_indices, + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + def read_state_pages_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + big_page_slot_ids: list[int], + layout: StateTransferLayout, + ) -> None: + """读入指定大页槽位。请求运行态由 restore_bounded_att_runtime_from_pages 另行恢复。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self.big_pages_pool + assert isinstance(big_page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" + # 本地目标保留本 TP rank 的分片形状。layout.regions 描述 CPU 全局字段。 + pool_page_fields = [ + self._tensor_fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids + ] + for copies in bounded_att_buffer_copies( + page_fields=pool_page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=cpu_page_indices, + write=False, + ): + copy_bounded_att_tensors(copies=copies) + + # ======================================================================== + # 四、DP 迁移：大页快照与基准运行态 + # ======================================================================== + + def copy_big_pages_from_dp(self, src_page_ids: list[int], dst_page_ids: list[int], remote_dp_bounded_state) -> None: + """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" + assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" + assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + src_pool = remote_dp_bounded_state.big_pages_pool + dst_pool = self.big_pages_pool + assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" + # 大页是基准快照，与请求槽位无关。页池须在本进程注册，当前流访问完才能归还页面。 + conv_copies, ssm_copies = [], [] + for src_page, dst_page in zip(src_page_ids, dst_page_ids): + conv_copies.append((src_pool.conv_state[src_page], dst_pool.conv_state[dst_page])) + ssm_copies.append((src_pool.ssm_state[src_page], dst_pool.ssm_state[dst_page])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + + def copy_runtime_state_from_dp( + self, + src_req_indices: list[int], + dst_req_indices: list[int], + remote_dp_bounded_state: "ReplayKDAState", + ) -> None: + """复制已接受运行态到目标单槽，并清空目标 history。""" + assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" + assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + width = self.config.conv_width + conv_copies = [] + for src_req, dst_req in zip(src_req_indices, dst_req_indices): + self.clear_history(dst_req) + conv_copies.append( + ( + remote_dp_bounded_state._accepted_conv(src_req), + self.conv_state[:, dst_req, :, :width], + ) + ) + copy_bounded_att_tensors(copies=conv_copies) + remote_dp_bounded_state.export_ssm(src_req_indices, [self.ssm_state[:, dst_req] for dst_req in dst_req_indices]) + if self.mtp_state_index is not None: + for dst_req in dst_req_indices: + self.mtp_state_index[dst_req].fill_(0) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" + return self._transfer_regions() + + def write_to_pd_page( + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """写出已接受 Conv 窗口和物化 SSM，不改写源 history。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + accepted = self.get_accepted_state(req_idx) + fields = self._tensor_fields(accepted.conv, accepted.ssm) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """读入基准 Conv/SSM，并清空该请求的 history。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + self.clear_history(req_idx) + fields = self._tensor_fields( + self.conv_state[:, req_idx, :, : self.config.conv_width], + self.ssm_state[:, req_idx], + ) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=False, + ): + copy_bounded_att_tensors(copies=copies) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + # ======================================================================== + # 六、Replay history + # ======================================================================== + def finish_prefill(self, logical_layer_id: int, reqs: torch.Tensor) -> None: + """最后一层把已接受前缀整段压进 SSM。""" + + layers = () if self.config is None else self.config.logical_layer_ids + if layers and logical_layer_id == layers[-1]: + self.compress_prefix(reqs) + + def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None: + """最后一层记下这 1 个 token，再按容量决定是否整段压缩。""" + + layers = () if self.config is None else self.config.logical_layer_ids + if not layers or logical_layer_id != layers[-1]: + return + self.accept_single_token(reqs) + self.maybe_compress_prefix(reqs) + + def clear_history(self, req: int | torch.Tensor) -> None: + """清空复用或恢复请求的临时 history，保留 canonical SSM。""" + + if type(req) is int: + # 标量赋值会把 CPU tensor 复制到 CUDA slot。 + self.cursors[req].zero_() + else: + self.cursors[req] = 0 + + def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: + """返回本步写入前的前缀长度。decode 不在这里修改 SSM。""" + + del cu_seqlens + return self.cursors.index_select(0, reqs.long()) + + def accept_updates( + self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None, width: int | None = None + ) -> None: + """更新 cursor。width 缺省时按 verify_width 取已接受长度，传 1 时固定加一个 token。""" + + _advance_cursor[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + accepted_state_indices, + reqs.numel(), + self.hold_req_id, + self.verify_width if width is None else width, + 256, + ) + + def accept_single_token(self, reqs: torch.Tensor) -> None: + """把一个已接受 token 记入前缀长度。""" + + if reqs.numel() == 0: + return + # width=1 不读 index。仍传 reqs，kernel 不能收到空指针。 + self.accept_updates(reqs, reqs, width=1) + + def merge_accepted_updates(self, reqs: torch.Tensor) -> None: + """把 cursor 覆盖的已接受前缀整段折回 canonical SSM，并清空 cursor。""" + + if reqs.numel() == 0: + return + self._materialize_accepted(reqs, self.ssm_state, snapshot=False) + self.cursors.index_fill_(0, reqs.long(), 0) + + def compress_prefix(self, reqs: torch.Tensor) -> None: + """把 cursor 覆盖的已接受前缀整段压进 SSM。长度为 0 时 kernel 直接返回。""" + + if reqs.numel() == 0: + return + if self.ssm_state.is_cuda: + self.merge_accepted_updates(reqs) + else: + self.clear_history(reqs) + + def maybe_compress_prefix(self, reqs: torch.Tensor) -> None: + """长度加上一次最大 verify_width 放不下，或 force_compress 打开时，整段压缩。""" + + if reqs.numel() == 0 or not self.ssm_state.is_cuda: + return + if self.force_compress: + self.compress_prefix(reqs) + return + self._materialize_accepted(reqs, self.ssm_state, snapshot=False, when_full=True) + + def accepted_ssm(self, req_idx: int) -> torch.Tensor: + """返回不修改 canonical SSM 和 history 的单请求快照。""" + + layers, _, hv, k, v = self.ssm_state.shape + output = torch.empty((layers, hv, k, v), dtype=self.ssm_state.dtype, device=self.ssm_state.device) + reqs = torch.tensor([req_idx], dtype=torch.int32, device=self.ssm_state.device) + self._materialize_accepted(reqs, output, snapshot=True) + return output + + def export_accepted_state( + self, + req_indices: list[int], + targets: list[torch.Tensor], + flush: bool = False, + ) -> None: + """将多个 accepted 快照直接写入调用方目标页。""" + + if not req_indices: + return + shape = (self.ssm_state.shape[0], *self.ssm_state.shape[2:]) + strides = targets[0].stride() + assert len(req_indices) == len(targets) + assert all(t.shape == shape and t.stride() == strides and t.dtype == self.ssm_state.dtype for t in targets) + device = self.ssm_state.device + self._materialize_accepted( + torch.tensor(req_indices, dtype=torch.int32, pin_memory=True).to(device, non_blocking=True), + torch.tensor([t.data_ptr() for t in targets], dtype=torch.int64, pin_memory=True).to( + device, non_blocking=True + ), + snapshot=not flush, + flush_export=flush, + output_strides=strides, + ) + + def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: + """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" + + if self.verify_width > 1: + commit_ssm_prefix(indices, reqs, lengths, self.cursors) + else: + commit_ssm_prefix(indices, reqs, lengths) + self.accept_updates(reqs, indices) + + def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: bool = False) -> None: + """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先物化再拷贝。""" + + if not req_indices: + return + target = targets[0] + if target.is_cuda and target.device == self.ssm_state.device: + self.export_accepted_state(req_indices, targets, flush=flush) + return + for req, dst in zip(req_indices, targets): + snap = self.accepted_ssm(req) + dst.copy_(snap) + if flush: + self.ssm_state[:, req].copy_(snap) + self.clear_history(req) + + _run_config: dict[str, int] | None = None + + @property + def run_config(self) -> dict[str, int] | None: + """返回物化和 autotune 使用的 Triton 配置，不是 SSM 布局。""" + + return self._run_config + + @run_config.setter + def run_config(self, config: dict | None) -> None: + """保存一份与调用方脱开的 kernel 配置。""" + + if config is None: + self._run_config = None + return + self._run_config = { + "BV": int(config["BV"]), + "num_warps": int(config["num_warps"]), + "num_stages": int(config.get("num_stages", 3)), + } + + # ======================================================================== + # 七、私有辅助实现 + # ======================================================================== + + def _transfer_regions(self) -> tuple[TransferRegion, ...]: + """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" + # 只传有效历史窗口，不包含 MTP 候选区和游标。 + config = self.config + storage_layer_count = len(config.logical_layer_ids) + conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim + conv_shape = (storage_layer_count, conv_channels, config.conv_width) + conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, Conv shape=%s", + self.state_name, + config.tp_rank, + config.tp_world_size, + conv_shape, + ) + + # Conv 按 Q/K/V 拼接。三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 + local_channel_start = 0 + global_channel_start = 0 + for name, local_head_count, global_head_count, head_dim in ( + ("Q", config.num_k_heads, config.global_k_heads, config.key_dim), + ("K", config.num_k_heads, config.global_k_heads, config.key_dim), + ("V", config.num_v_heads, config.global_v_heads, config.value_dim), + ): + # 每组连续 rank 保存相同 heads，仅组内首 rank 写出。 + replica_count = config.tp_world_size * local_head_count // global_head_count + head_start = config.tp_rank // replica_count * local_head_count + write_enabled = config.tp_rank % replica_count == 0 + local_channel_count = local_head_count * head_dim + target_channel_start = global_channel_start + head_start * head_dim + regions.append( + TransferRegion( + field="conv", + byte_start=0, + byte_end=conv_byte_end, + dtype=config.conv_dtype, + axes=("layer", "channel", "window"), + shape=conv_shape, + axis_mappings=( + AxisMapping( + axis="channel", + local_start=local_channel_start, + external_start=target_channel_start, + length=local_channel_count, + ), + AxisMapping(axis="window", local_start=0, external_start=0, length=config.conv_width), + ), + write_enabled=write_enabled, + ) + ) + logger.debug( + "%s Conv %s: local_channel_start=%d, global_channel_start=%d, " + "channels=%d, replicas=%d, write_enabled=%s", + self.state_name, + name, + local_channel_start, + target_channel_start, + local_channel_count, + replica_count, + write_enabled, + ) + local_channel_start += local_channel_count + global_channel_start += global_head_count * head_dim + + # SSM 沿 V head 切分，接在 Conv 之后，并按 SSM dtype 对齐字节起点。 + replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads + head_start = config.tp_rank // replica_count * config.num_v_heads + write_enabled = config.tp_rank % replica_count == 0 + ssm_shape = (storage_layer_count, config.global_v_heads, config.key_dim, config.value_dim) + element_bytes = config.ssm_dtype.itemsize + ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes + ssm_region = TransferRegion( + field="ssm", + byte_start=ssm_byte_start, + byte_end=ssm_byte_start + prod(ssm_shape) * element_bytes, + dtype=config.ssm_dtype, + axes=("layer", "head", "key_dim", "value_dim"), + shape=ssm_shape, + axis_mappings=( + AxisMapping(axis="head", local_start=0, external_start=head_start, length=config.num_v_heads), + AxisMapping(axis="key_dim", local_start=0, external_start=0, length=config.key_dim), + AxisMapping(axis="value_dim", local_start=0, external_start=0, length=config.value_dim), + ), + write_enabled=write_enabled, + ) + regions.append(ssm_region) + logger.debug( + "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", + self.state_name, + head_start, + config.num_v_heads, + write_enabled, + ssm_region.byte_start, + ssm_region.byte_end, + ) + return tuple(regions) + + def _pool(self, page_kind: PageKind): + """根据大小页类型取得已启用的页池，不访问 allocator。""" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" + return pool + + @staticmethod + def _tensor_fields(conv: torch.Tensor, ssm: torch.Tensor): + """描述一页里的 Conv 和 SSM 字段。""" + return { + "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), + "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), + } + + def _accepted_conv(self, req_idx: int) -> torch.Tensor: + """取出该请求已接受的 Conv 窗口，不包含 MTP 候选区。""" + width = self.config.conv_width + conv = self.conv_state[:, req_idx] + if self.mtp_state_index is None: + return conv[:, :, :width] + offset = self.mtp_state_index[req_idx] + return conv.index_select(-1, torch.arange(width, device=conv.device) + offset) + + def _materialize_accepted( + self, + reqs: torch.Tensor, + output: torch.Tensor, + snapshot: bool, + output_strides=None, + flush_export: bool = False, + when_full: bool = False, + ) -> None: + """将请求的 accepted prefix 写入 canonical SSM 或调用方提供的输出页。""" + + assert not (snapshot and flush_export), "flush export must update the canonical state" + layers, slots, hv, k, v = self.ssm_state.shape + config = self.run_config or {"BV": 32, "num_warps": 4, "num_stages": 3} + bv = config["BV"] + _fold_prefix[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( + self.ssm_state, + output, + self.raw_keys, + self.raw_values, + self.gates, + self.betas, + self.cursors, + reqs, + self.ssm_state.stride(), + slots, + hv, + self.num_key_heads, + k, + v, + self.capacity, + self.hold_req_id, + triton.next_power_of_2(k), + bv, + snapshot, + True, + FLUSH_EXPORT=flush_export, + OUTPUT_STRIDES=output_strides, + WHEN_FULL=when_full, + NEXT=self.verify_width, + num_warps=config["num_warps"], + num_stages=config["num_stages"], + ) + if flush_export or when_full: + _reset_folded_cursors[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + reqs.numel(), + self.hold_req_id, + 0 if flush_export else self.capacity, + self.verify_width, + 256, + ) + + +def _history_owner( + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype = torch.bfloat16, + num_key_heads: int | None = None, + gate_lower_bound: float = -5.0, + compute_mode: str = "inline", + run_config: dict[str, int] | None = None, +) -> ReplayKDAState: + """把 history 挂到调用方已有的 SSM 上，不再分配第二份状态。供保持原 stride 的 kernel 测试使用。""" + + owner = ReplayKDAState.__new__(ReplayKDAState) + owner.config = None + owner.state_name = None + owner.request_capacity = ssm.shape[1] + owner.mtp_step = verify_width - 1 + owner._algorithm = "kda" + owner._logical_to_storage_layer = {} + owner.conv_state = None + owner.mtp_state_index = None + owner.big_pages_pool = None + owner.small_pages_pool = None + _bind_history( + owner, + ssm, + capacity, + verify_width, + activation_dtype, + num_key_heads, + gate_lower_bound, + compute_mode, + run_config, + ) + return owner + + +def _bind_history( + owner, + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype, + num_key_heads: int | None, + gate_lower_bound: float, + compute_mode: str, + run_config: dict[str, int] | None, + force_compress: bool = False, +) -> None: + """在 owner 上创建 KDA replay 的连续前缀 history 和 cursor。""" + + assert ssm.dtype in (torch.float32, torch.bfloat16), "KDA replay requires FP32 or BF16 SSM state" + assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 + assert compute_mode in ("inline", "precompute") + owner.ssm_state = ssm + owner.kda = True + owner.gate_lower_bound = gate_lower_bound + owner.capacity = capacity + owner.verify_width = verify_width + owner.force_compress = bool(force_compress) + owner.compute_mode = compute_mode + owner.run_config = run_config + owner._config_is_fixed = run_config is not None + owner.batch_configs = {} + layers, slots, hv, k, v = ssm.shape + owner.num_key_heads = hv if num_key_heads is None else num_key_heads + assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 + owner.hold_req_id = slots - 1 + owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) + # 归一化 k 可用 fp16；delta 保留 fp32，避免有限 BF16 输入的修正量溢出 fp16。 + # 原始 k/v 保持激活精度。g 和 beta 保持 fp32。压缩在 fp32 里从原始值折回。 + reconstruction_dtype = torch.float16 if activation_dtype == torch.bfloat16 else torch.float32 + key_shape = (layers, slots, owner.num_key_heads, capacity, k) + owner.keys = torch.empty(key_shape, dtype=reconstruction_dtype, device=ssm.device) + owner.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=torch.float32, device=ssm.device) + owner.gates = torch.empty((layers, slots, hv, capacity, k), dtype=torch.float32, device=ssm.device) + owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) + owner.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=activation_dtype, device=ssm.device) + owner.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) diff --git a/lightllm/common/attention_state/bounded_att/utils.py b/lightllm/common/attention_state/bounded_att/utils.py index 93f5b2ebe3..be4be7e5ed 100644 --- a/lightllm/common/attention_state/bounded_att/utils.py +++ b/lightllm/common/attention_state/bounded_att/utils.py @@ -2,7 +2,8 @@ from __future__ import annotations -from math import prod +from math import lcm, prod +from mmap import PAGESIZE from weakref import finalize import torch @@ -17,6 +18,10 @@ CUDA_HOST_REGISTER_PORTABLE = 1 # Mapped：建立设备可访问映射，供 CUDA/Triton kernel 直接访问主机存储。 CUDA_HOST_REGISTER_MAPPED = 2 +# Keep DP-replicated shared page pools below the CUDA host-registration range +# pressure observed with eight ranks. Each chunk is still rounded to a whole +# snapshot, so DMA copies never cross registration boundaries. +_HOST_REGISTRATION_BYTES = 1024 * 1024 * 1024 # ======================================================================== @@ -43,16 +48,25 @@ def create_page_pool_tensor( # 阶段一：未启用 DP fetch 时沿用 torch 分配；GPU 页池也保持原设备分配方式。 if not shared: - tensor = torch.empty( - size=size, - dtype=dtype, - device=device, - pin_memory=pin_memory, - ) + try: + tensor = torch.empty( + size=size, + dtype=dtype, + device=device, + pin_memory=pin_memory, + ) + except RuntimeError: + if device.type != "cpu" or not pin_memory or not torch.cuda.is_available(): + raise + tensor = torch.empty(size=size, dtype=dtype, device=device, pin_memory=False) + register_page_pool_tensor(tensor=tensor) + allocation = "torch.empty + cudaHostRegister" + else: + allocation = "torch.empty" reason = "GPU device allocation" if device.type != "cpu" else "DP prompt cache fetch disabled" logger.info( f"Allocated bounded attention page pool: shape={size}, dtype={dtype}, bytes={nbytes}, " - f"device={device}, shared=False, pin_memory={pin_memory}, allocator=torch.empty, reason={reason}" + f"device={device}, shared=False, pin_memory={pin_memory}, allocator={allocation}, reason={reason}" ) return tensor @@ -91,24 +105,43 @@ def register_page_pool_tensor(tensor: torch.Tensor) -> None: """ if tensor.device.type != "cpu" or tensor.numel() == 0 or tensor.is_pinned(): return - assert tensor.is_shared(), "page pool registration requires shared CPU storage" storage = tensor.untyped_storage() ptr = storage.data_ptr() nbytes = storage.nbytes() device_id = torch.cuda.current_device() flags = CUDA_HOST_REGISTER_PORTABLE | CUDA_HOST_REGISTER_MAPPED - torch.cuda.check_error(torch.cuda.cudart().cudaHostRegister(ptr, nbytes, flags)) + # 大页池整块注册在多 GPU 初始化时可能失败；边界对齐完整快照，避免一次 DMA 跨注册区间。 + registration_unit = lcm(prod(tensor.shape[1:]) * tensor.element_size(), PAGESIZE) + chunk_bytes = max(registration_unit, _HOST_REGISTRATION_BYTES // registration_unit * registration_unit) + if not tensor.is_shared(): + chunk_bytes = nbytes + addresses = [] + runtime = torch.cuda.cudart() + try: + for offset in range(0, nbytes, chunk_bytes): + address = ptr + offset + torch.cuda.check_error(runtime.cudaHostRegister(address, min(chunk_bytes, nbytes - offset), flags)) + addresses.append(address) + except BaseException: + logger.error( + f"Failed to register bounded attention page pool: bytes={nbytes}, cuda_device={device_id}, " + f"chunk_bytes={chunk_bytes}, registered_chunks={len(addresses)}" + ) + if addresses: + _unregister_page_pool_storage(tuple(addresses), device_id) + raise # finalizer 绑定 Storage，所有本地 Tensor 及视图释放后才注销，不能绑在单个 Tensor 上。 # 正常释放时等待 GPU 完成使用；解释器退出时交给进程清理，避免 CUDA 已关闭后再调用。 cleanup = finalize( storage, _unregister_page_pool_storage, - ptr, + tuple(addresses), device_id, ) cleanup.atexit = False logger.info( - f"Registered shared bounded attention page pool: bytes={nbytes}, cuda_device={device_id}, " + f"Registered shared bounded attention page pool: bytes={nbytes}, chunks={len(addresses)}, " + f"cuda_device={device_id}, " "flags=Portable|Mapped; registration applies to this process's mapping" ) @@ -118,8 +151,9 @@ def register_page_pool_tensor(tensor: torch.Tensor) -> None: # ======================================================================== -def _unregister_page_pool_storage(ptr: int, device_id: int) -> None: +def _unregister_page_pool_storage(addresses: tuple[int, ...], device_id: int) -> None: """等待当前 GPU 使用结束，再注销本进程的映射；共享内存本身由 Storage 管理。""" torch.cuda.synchronize(device=device_id) - torch.cuda.check_error(torch.cuda.cudart().cudaHostUnregister(ptr)) - logger.debug(f"Unregistered shared bounded attention page pool: cuda_device={device_id}, ptr={ptr}") + for ptr in reversed(addresses): + torch.cuda.check_error(torch.cuda.cudart().cudaHostUnregister(ptr)) + logger.debug(f"Unregistered shared bounded attention page pool: cuda_device={device_id}, chunks={len(addresses)}") 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..c658533abc 100644 --- a/lightllm/common/attention_state/full_att/gqa/gqa_state.py +++ b/lightllm/common/attention_state/full_att/gqa/gqa_state.py @@ -146,14 +146,13 @@ def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, if not len(token_indices): 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 + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv_fused - # 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] + destindex_copy_kv_fused( + KV=kv, + DestLoc=token_indices, + OutK=self.k_storage[storage_layer_index], + OutV=self.v_storage[storage_layer_index], ) def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: @@ -169,22 +168,16 @@ def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Ten 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 + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_pair_to_kv_buffer_pair # 直接使用调用方提供的 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( + copy_kv_buffer_pair_to_kv_buffer_pair( src_mem_index=src_indices, dst_mem_index=dst_indices, - kv_buffer=self.v_storage, + k_buffer=self.k_storage, + v_buffer=self.v_storage, ) # ======================================================================== diff --git a/lightllm/common/basemodel/attention/create_linear_utils.py b/lightllm/common/basemodel/attention/create_linear_utils.py index 6a40a05183..fce8e8d78f 100644 --- a/lightllm/common/basemodel/attention/create_linear_utils.py +++ b/lightllm/common/basemodel/attention/create_linear_utils.py @@ -19,38 +19,48 @@ def get_qwen35_linear_prefill_att_backend_class(index=1, priority_list=("flashql args = get_env_start_args() backend_str = _get_backend_str(args.llm_prefill_att_backend, index) if backend_str != "auto": - return load_backend_class(linear_prefill_att_backend_classes[backend_str]) + return _linear_backend_class(linear_prefill_att_backend_classes[backend_str], backend_str) for backend_name in priority_list: backend_class = linear_prefill_att_backend_classes[backend_name] if backend_name == "triton": logger.info("Linear prefill attention backend: triton.") - return load_backend_class(backend_class) + return _linear_backend_class(backend_class, backend_name) if validate(backend_name): logger.info(f"Linear prefill attention backend: {backend_name} (validated).") - return load_backend_class(backend_class) + return _linear_backend_class(backend_class, backend_name) logger.warning("No linear prefill attention backend validation succeeded, falling back to Triton.") - return load_backend_class("linear.triton:TritonLinearAttBackend") + return _linear_backend_class("linear.triton:TritonLinearAttBackend", "triton") def get_qwen35_linear_decode_att_backend_class(index=1, priority_list=("triton",)): args = get_env_start_args() backend_str = _get_backend_str(args.llm_decode_att_backend, index) if backend_str != "auto": - return load_backend_class(linear_decode_att_backend_classes[backend_str]) + return _linear_backend_class(linear_decode_att_backend_classes[backend_str], backend_str) for backend_name in priority_list: backend_class = linear_decode_att_backend_classes[backend_name] if backend_name == "triton": logger.info("Linear decode attention backend: triton.") - return load_backend_class(backend_class) + return _linear_backend_class(backend_class, backend_name) if validate(backend_name): logger.info(f"Linear decode attention backend: {backend_name} (validated).") - return load_backend_class(backend_class) + return _linear_backend_class(backend_class, backend_name) logger.warning("No linear decode attention backend validation succeeded, falling back to Triton.") - return load_backend_class("linear.triton:TritonLinearAttBackend") + return _linear_backend_class("linear.triton:TritonLinearAttBackend", "triton") + + +def _linear_backend_class(native_path: str, backend_name: str): + """Native 仍按路径加载。replay 把 kernel 名字交给独立的 GDN replay 入口。""" + + if getattr(get_env_start_args(), "ssm_state_mode", "native") != "replay": + return load_backend_class(native_path) + from lightllm.common.basemodel.attention.linear.replay_gdn import replay_gdn_backend_class + + return replay_gdn_backend_class(backend_name) def _get_backend_str(backend_args, index): diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 986214568f..48e3eb64d4 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -8,7 +8,6 @@ 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, ) @@ -17,6 +16,7 @@ ) from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( fused_recurrent_gated_delta_rule, + fused_recurrent_gated_delta_rule_packed_decode, ) if TYPE_CHECKING: @@ -37,7 +37,6 @@ def get_prefill_kernel(self): pass def _init_linear_layer_metadata(self, network_config, tp_world_size): - self.mtp_step = get_env_start_args().mtp_step # Linear attention specific dimensions @@ -109,7 +108,6 @@ def create_att_decode_state(self, infer_state: "InferStateInfo") -> "LinearAttDe @dataclasses.dataclass class LinearAttPrefillAttState(BasePrefillAttState): - b_conv_buffer_idx: torch.Tensor = None b_ssm_buffer_idx: torch.Tensor = None @@ -180,6 +178,7 @@ def _gdn_prefill_kernel( has_initial_state=infer_state.b_ready_cache_len > 0, conv_states=conv_states, activation=backend.activation, + max_seqlen=infer_state.max_q_seq_len, ) mixed_qkv = out_tensor.transpose(0, 1) @@ -207,11 +206,11 @@ def _gdn_prefill_kernel( @dataclasses.dataclass class LinearAttDecodeAttState(BaseDecodeAttState): - b_conv_buffer_idx: torch.Tensor = None b_ssm_buffer_idx: torch.Tensor = None b1_mtp_cu_q_seq_len: torch.Tensor = None b_num_accepted_tokens: torch.Tensor = None + b1_mtp_cu_q_seq_len_int64: 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) @@ -232,6 +231,7 @@ def init_state(self): 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) + self.b1_mtp_cu_q_seq_len_int64 = self.b1_mtp_cu_q_seq_len.to(torch.int64) def _init_normal_decode_state(self): self.b_conv_buffer_idx = self.infer_state.b_req_idx @@ -345,37 +345,30 @@ def _gdn_decode_kernel( ): 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( + from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_update + + # Keep the post-convolution q/k/v layout packed. The recurrent kernel + # performs the GQA rearrangement and recurrent update in one launch, + # matching the packed decode path used by vLLM and avoiding six + # short-lived q/k/v/z/a/b tensors on every linear-attention layer. + mixed_qkv = causal_conv1d_update( 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, + activation=backend.activation, + conv_state_indices=self.b_conv_buffer_idx, ) - core_attn_out, _ = fused_recurrent_gated_delta_rule( - q=query, - k=key, - v=value, + core_attn_out, _ = fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv=mixed_qkv, + a=a, + b=b, + A_log=layer_weight.linear_A_log.weight, + dt_bias=layer_weight.linear_dt_bias.weight, + scale=backend.head_k_dim ** -0.5, 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 @@ -418,7 +411,7 @@ def _gdn_mtp_kernel( k=key, v=value, initial_state=ssm_states, - cu_seqlens=cu_seqlens_q.to(torch.long), + cu_seqlens=self.b1_mtp_cu_q_seq_len_int64, 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, @@ -426,5 +419,10 @@ def _gdn_mtp_kernel( dt_bias=layer_weight.linear_dt_bias.weight, a_raw=a, b_raw=b, + run_config=( + {"BV": 32, "num_warps": 4, "num_stages": 1} + if getattr(backend.model, "linear_att_mtp_match_decode", False) + else None + ), ) return core_attn_out diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py new file mode 100644 index 0000000000..d08a714d27 --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -0,0 +1,352 @@ +"""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 + + +def replay_gdn_backend_class(name: str): + """按已选中的 linear kernel 名字返回对应的 replay backend。""" + + classes = { + "triton": ReplayGDNTritonAttBackend, + "flashqla": ReplayGDNFlashQlaAttBackend, + } + try: + return classes[name] + except KeyError: + raise ValueError(f"replay GDN has no backend named {name!r}") from None + + +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) + state.finish_prefill(layer_num, self.b_conv_buffer_idx) + 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_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) + state.finish_decode_token(layer_num, self.b_conv_buffer_idx) + 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 + 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_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_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, + ) + 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.b1_mtp_cu_q_seq_len, + ) + 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 = _rearrange_mixed_qkv(backend, mixed_qkv, decode=False) + 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.b1_mtp_cu_q_seq_len, + ) + 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_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py new file mode 100644 index 0000000000..71fcf8e1aa --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -0,0 +1,86 @@ +"""GDN replay decode 的 kernel 辅助。 + +History 张量和 cursor 在 ReplayGDNState。这里只选择 kernel 配置，并发起一层重放。 +""" + +from __future__ import annotations + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import _decode_replay +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config + + +def replay( + history, + linear_layer_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + run_config: dict[str, int] | None = None, +) -> tuple[torch.Tensor, dict[str, int]]: + """执行一层 GDN replay decode。Kernel 只接收这一层的 SSM 和 history；state 只提供共用 cursor 与配置。""" + + if history.kda: + raise ValueError("GDN replay step does not use the KDA gate") + config = run_config or resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + hv, kd, vd = history.ssm_state.shape[-3:] + assert history.num_key_heads in (q.shape[-2], hv) + axis = 1 if cu_seqlens is not None else 0 + out = torch.empty_like(v) + bv = config["BV"] + # mtp_step == 0 时 verify_width 为 1。compute_mode=precompute 仍然投影；BV < 16 才退回逐 token。 + precompute_projections = history.compute_mode == "precompute" and bv >= 16 + # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 + bw = triton.next_power_of_2(history.verify_width) + _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( + q, + k, + v, + a, + b, + a_log, + bias, + history.ssm_state[linear_layer_idx], + history.keys[linear_layer_idx], + history.deltas[linear_layer_idx], + history.gates[linear_layer_idx], + history.raw_keys[linear_layer_idx], + history.raw_values[linear_layer_idx], + history.betas[linear_layer_idx], + reqs, + history.cursors, + cu_seqlens, + out, + q.stride(axis), + k.stride(axis), + v.stride(axis), + a.stride(0), + b.stride(0), + q.shape[-2], + hv, + history.num_key_heads, + kd, + vd, + history.capacity, + history.hold_req_id, + triton.next_power_of_2(kd), + bv, + cu_seqlens is not None, + precompute_projections, + bw, + False, + history.gate_lower_bound, + BQK=max(16, 2 * bw), + BR=triton.next_power_of_2(max(16, history.capacity + bw)), + num_warps=config["num_warps"], + num_stages=config.get("num_stages", 3), + ) + return out, config diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py new file mode 100644 index 0000000000..3090faa9a8 --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -0,0 +1,241 @@ +"""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) + state.finish_prefill(layer_num, self.b_conv_buffer_idx) + 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_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) + 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 = params["raw_gate"].view(*shape, backend.tp_hidden_size) + raw_beta = params["raw_beta"].view(*shape, backend.tp_num_heads) + output, _ = replay( + state, + state.storage_layer(layer_num), + q, + k, + v, + raw_gate.reshape(-1, backend.tp_hidden_size), + raw_beta.reshape(-1, backend.tp_num_heads), + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, + self.b_conv_buffer_idx, + self.b1_mtp_cu_q_seq_len, + ) + if backend.mtp_step == 0: + state.finish_decode_token(layer_num, self.b_conv_buffer_idx) + 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 + 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_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_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/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py new file mode 100644 index 0000000000..55514aeb57 --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -0,0 +1,86 @@ +"""KDA replay decode 的 kernel 辅助。 + +History 张量和 cursor 在 ReplayKDAState。这里只选择 kernel 配置，并发起一层重放。 +""" + +from __future__ import annotations + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import _decode_replay +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config + + +def replay( + history, + linear_layer_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + run_config: dict[str, int] | None = None, +) -> tuple[torch.Tensor, dict[str, int]]: + """执行一层 KDA replay decode。Kernel 只接收这一层的 SSM 和 history；state 只提供共用 cursor 与配置。""" + + if not history.kda: + raise ValueError("KDA replay step requires the KDA gate") + config = run_config or resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + hv, kd, vd = history.ssm_state.shape[-3:] + assert history.num_key_heads in (q.shape[-2], hv) + axis = 1 if cu_seqlens is not None else 0 + out = torch.empty_like(v) + bv = config["BV"] + # mtp_step == 0 时 verify_width 为 1。compute_mode=precompute 仍然投影；BV < 16 才退回逐 token。 + precompute_projections = history.compute_mode == "precompute" and bv >= 16 + # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 + bw = triton.next_power_of_2(history.verify_width) + _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( + q, + k, + v, + a, + b, + a_log, + bias, + history.ssm_state[linear_layer_idx], + history.keys[linear_layer_idx], + history.deltas[linear_layer_idx], + history.gates[linear_layer_idx], + history.raw_keys[linear_layer_idx], + history.raw_values[linear_layer_idx], + history.betas[linear_layer_idx], + reqs, + history.cursors, + cu_seqlens, + out, + q.stride(axis), + k.stride(axis), + v.stride(axis), + a.stride(0), + b.stride(0), + q.shape[-2], + hv, + history.num_key_heads, + kd, + vd, + history.capacity, + history.hold_req_id, + triton.next_power_of_2(kd), + bv, + cu_seqlens is not None, + precompute_projections, + bw, + True, + history.gate_lower_bound, + BQK=max(16, 2 * bw), + BR=triton.next_power_of_2(max(16, history.capacity + bw)), + num_warps=config["num_warps"], + num_stages=config.get("num_stages", 3), + ) + return out, config diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/pinned_embedding_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/pinned_embedding_weight.py new file mode 100644 index 0000000000..3cb376594a --- /dev/null +++ b/lightllm/common/basemodel/layer_weights/meta_weights/pinned_embedding_weight.py @@ -0,0 +1,66 @@ +import torch + +from .row_sharded_embedding_weight import RowShardedEmbeddingWeight +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +class PinnedRowShardedEmbeddingWeight(RowShardedEmbeddingWeight): + def __init__( + self, dim, vocab_size, weight_name, data_type, shard_count, *, cpu_allocator, tp_rank=None, tp_world_size=None + ): + if not callable(cpu_allocator): + raise ValueError("Pinned embeddings require a CPU allocation policy") + self.cpu_allocator = cpu_allocator + super().__init__(dim, vocab_size, weight_name, data_type, shard_count, tp_rank, tp_world_size) + + def _create_weight(self): + shape = (self.tp_vocab_end_id - self.tp_vocab_start_id, self.dim) + with torch.cuda.device(self.device_id_): + self.weight = self.cpu_allocator(shape, dtype=self.data_type_, device="cpu", pin_memory=True) + if ( + self.weight.device.type != "cpu" + or self.weight.dtype != self.data_type_ + or tuple(self.weight.shape) != shape + or not self.weight.is_contiguous() + or (self.weight.numel() and not self.weight.is_pinned()) + ): + raise ValueError( + "Embedding allocator must return contiguous pinned CPU storage with the requested shape/dtype" + ) + self.weight.load_ok = False + logger.info( + f"Pinned embedding table: rank={self.tp_rank_} rows={shape[0]} dim={self.dim} " + f"bytes={self.weight.numel() * self.weight.element_size()} dtype={self.weight.dtype} residency=cpu" + ) + + def load_hf_weights(self, weights): + names = {self.weight_name} | {f"{self._shard_prefix}.shard_{index}.weight" for index in self._required_shards} + host_weights = {name: value.cpu() if name in names else value for name, value in weights.items()} + super().load_hf_weights(host_weights) + + def _cuda_forward(self, input_ids, out=None, alloc_func=torch.empty): + if ( + not input_ids.is_cuda + or input_ids.ndim != 1 + or not input_ids.is_contiguous() + or input_ids.dtype not in (torch.int32, torch.int64) + ): + raise ValueError("Pinned embedding lookup requires a CUDA vector of integer IDs") + shape = (input_ids.shape[0], self.dim) + if out is None: + out = alloc_func(shape, dtype=self.weight.dtype, device=input_ids.device) + if ( + tuple(out.shape) != shape + or out.dtype != self.weight.dtype + or out.device != input_ids.device + or out.stride(1) != 1 + ): + raise ValueError("Pinned embedding output must match CUDA IDs, table dtype and lookup shape") + if not input_ids.numel(): + return out + return super()._triton_forward(input_ids, out=out, alloc_func=alloc_func) + + def _native_forward(self, *args, **kwargs): + raise ValueError("Pinned host embedding lookup requires the CUDA Triton implementation") diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/row_sharded_embedding_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/row_sharded_embedding_weight.py new file mode 100644 index 0000000000..ca2390c736 --- /dev/null +++ b/lightllm/common/basemodel/layer_weights/meta_weights/row_sharded_embedding_weight.py @@ -0,0 +1,64 @@ +from typing import Dict, Optional + +import torch + +from .embedding_weight import EmbeddingWeight + + +class RowShardedEmbeddingWeight(EmbeddingWeight): + def __init__( + self, + dim: int, + vocab_size: int, + weight_name: str, + data_type: torch.dtype, + shard_count: int, + tp_rank: Optional[int] = None, + tp_world_size: Optional[int] = None, + ): + if shard_count < 1 or vocab_size < 1 or vocab_size % shard_count: + raise ValueError("Row-sharded embeddings require a positive vocabulary divisible by shard count") + if not weight_name.endswith(".weight"): + raise ValueError("Embedding weight name must end with .weight") + self.shard_count = shard_count + self.shard_rows = vocab_size // shard_count + self._shard_prefix = weight_name[: -len(".weight")] + self._loaded_shards = set() + self._full_loaded = False + super().__init__(dim, vocab_size, weight_name, data_type, tp_rank, tp_world_size) + self._required_shards = { + index + for index in range(shard_count) + if max(index * self.shard_rows, self.tp_vocab_start_id) + < min((index + 1) * self.shard_rows, self.tp_vocab_end_id) + } + self.weight.load_ok = not self._required_shards + + def load_hf_weights(self, weights: Dict[str, torch.Tensor]): + shard_keys = [(index, f"{self._shard_prefix}.shard_{index}.weight") for index in range(self.shard_count)] + present = [(index, key) for index, key in shard_keys if key in weights] + if self.weight_name in weights: + if present or self._loaded_shards: + raise ValueError("Cannot mix full and row-sharded embedding checkpoint tensors") + if weights[self.weight_name].shape != (self.vocab_size, self.dim): + raise ValueError("Full embedding checkpoint dimensions do not match") + super().load_hf_weights(weights) + self._full_loaded = True + return + if self._full_loaded: + if present: + raise ValueError("Cannot mix full and row-sharded embedding checkpoint tensors") + return + for index, key in present: + source = weights[key] + if source.shape != (self.shard_rows, self.dim): + raise ValueError(f"Unexpected dimensions for embedding shard {index}: {tuple(source.shape)}") + shard_start = index * self.shard_rows + begin = max(shard_start, self.tp_vocab_start_id) + end = min(shard_start + self.shard_rows, self.tp_vocab_end_id) + if begin < end: + self.weight[begin - self.tp_vocab_start_id : end - self.tp_vocab_start_id].copy_( + source[begin - shard_start : end - shard_start].to(self.data_type_) + ) + self._loaded_shards.add(index) + self.weight.load_ok = self._loaded_shards >= self._required_shards diff --git a/lightllm/common/basemodel/multimodal_tokenizer.py b/lightllm/common/basemodel/multimodal_tokenizer.py index cdcbd7f089..6cd0f8f26c 100644 --- a/lightllm/common/basemodel/multimodal_tokenizer.py +++ b/lightllm/common/basemodel/multimodal_tokenizer.py @@ -40,6 +40,9 @@ def __getattr__(self, name): return obj_dict[name] return getattr(self.tokenizer, name) + def count_prompt_tokens(self, text_tokens, image_tokens, image_count, audio_tokens, audio_count): + return text_tokens + image_tokens + image_count + audio_tokens + audio_count + @abstractmethod def encode(self, prompt, multimodal_params: MultimodalParams = None, **kwargs): pass diff --git a/lightllm/common/basemodel/tp_reduce.py b/lightllm/common/basemodel/tp_reduce.py new file mode 100644 index 0000000000..f50c80b989 --- /dev/null +++ b/lightllm/common/basemodel/tp_reduce.py @@ -0,0 +1,58 @@ +import torch +import triton +import triton.language as tl + +from lightllm.distributed import all_gather_into_tensor + + +@triton.jit +def _sum_shards( + X, + Y, + N: tl.constexpr, + START: tl.constexpr, + SIZE: tl.constexpr, + WORLD: tl.constexpr, + BLOCK: tl.constexpr, +): + index = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + accumulator = tl.full((BLOCK,), 0, tl.float64) + for rank in range(WORLD): + value = tl.load(X + rank * N + START + index, index < SIZE, other=0) + accumulator += value.to(tl.float64) + tl.store(Y + index, accumulator, index < SIZE) + + +def ordered_tp_sum(input: torch.Tensor, group, world_size: int, *, scatter_rank=None) -> torch.Tensor: + if world_size == 1: + return input + if scatter_rank is not None: + if not 0 <= scatter_rank < world_size or input.shape[0] % world_size: + raise ValueError("Sequence-parallel input must divide evenly across valid ranks") + shape = (input.shape[0] // world_size, *input.shape[1:]) + output = torch.empty(shape, device=input.device, dtype=input.dtype) + start = scatter_rank * output.numel() + else: + output = input if input.is_contiguous() else torch.empty(input.shape, device=input.device, dtype=input.dtype) + start = 0 + if input.numel(): + source = input.contiguous() + symmetric = getattr(group, "symm_mem_reduce", None) if input.is_cuda else None + gathered = symmetric.ordered_all_gather(source) if symmetric is not None else None + if gathered is None: + gathered = torch.empty((world_size, *input.shape), device=input.device, dtype=input.dtype) + all_gather_into_tensor(gathered.view(-1, *input.shape[1:]), source, group=group) + if input.is_cuda: + _sum_shards[(triton.cdiv(output.numel(), 256),)]( + gathered, output, input.numel(), start, output.numel(), world_size, 256 + ) + else: + accumulator = torch.zeros(output.numel(), dtype=torch.float64, device=input.device) + shards = gathered.reshape(world_size, -1) + for rank in range(world_size): + accumulator.add_(shards[rank, start : start + output.numel()]) + output.copy_(accumulator.reshape(output.shape)) + if scatter_rank is None and output is not input: + input.copy_(output) + return input + return output diff --git a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py b/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py index 060a92bf7c..f301a03a1e 100644 --- a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py +++ b/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py @@ -34,6 +34,43 @@ def _fwd_kernel_destindex_copy_kv( return +@triton.jit +def _fwd_kernel_destindex_copy_kv_fused( + KV, + Dest_loc, + OutK, + OutV, + stride_kv_bs, + stride_kv_h, + stride_kv_d, + stride_ok_bs, + stride_ok_h, + stride_ok_d, + stride_ov_bs, + stride_ov_h, + stride_ov_d, + head_num, + head_dim, + BLOCK_DMODEL: tl.constexpr, + BLOCK_HEAD: tl.constexpr, +): + cur_index = tl.program_id(0) + offs_h = tl.arange(0, BLOCK_HEAD) + offs_d = (tl.arange(0, BLOCK_DMODEL)) % head_dim + + dest_index = tl.load(Dest_loc + cur_index).to(tl.int64) + mask = offs_h[:, None] < head_num + + kv_base = KV + cur_index * stride_kv_bs + offs_h[:, None] * stride_kv_h + offs_d[None, :] * stride_kv_d + v_base = KV + cur_index * stride_kv_bs + (head_num + offs_h[:, None]) * stride_kv_h + offs_d[None, :] * stride_kv_d + k_out = OutK + dest_index * stride_ok_bs + offs_h[:, None] * stride_ok_h + offs_d[None, :] * stride_ok_d + v_out = OutV + dest_index * stride_ov_bs + offs_h[:, None] * stride_ov_h + offs_d[None, :] * stride_ov_d + + tl.store(k_out, tl.load(kv_base, mask=mask, other=0.0), mask=mask) + tl.store(v_out, tl.load(v_base, mask=mask, other=0.0), mask=mask) + return + + @torch.no_grad() def destindex_copy_kv(K, DestLoc, Out): seq_len = DestLoc.shape[0] @@ -62,3 +99,40 @@ def destindex_copy_kv(K, DestLoc, Out): num_stages=1, ) return + + +@torch.no_grad() +def destindex_copy_kv_fused(KV, DestLoc, OutK, OutV): + """Copy packed K/V rows to separate storage buffers with one kernel launch.""" + seq_len = DestLoc.shape[0] + head_num = OutK.shape[1] + head_dim = OutK.shape[2] + assert KV.ndim == 3 and KV.shape[0] == seq_len + assert OutK.shape == OutV.shape + assert KV.shape[1] == 2 * head_num and KV.shape[2] == head_dim + + BLOCK_HEAD = triton.next_power_of_2(head_num) + grid = (seq_len,) + + _fwd_kernel_destindex_copy_kv_fused[grid]( + KV, + DestLoc, + OutK, + OutV, + KV.stride(0), + KV.stride(1), + KV.stride(2), + OutK.stride(0), + OutK.stride(1), + OutK.stride(2), + OutV.stride(0), + OutV.stride(1), + OutV.stride(2), + head_num, + head_dim, + BLOCK_DMODEL=triton.next_power_of_2(head_dim), + BLOCK_HEAD=BLOCK_HEAD, + num_warps=1, + num_stages=1, + ) + return diff --git a/lightllm/common/basemodel/triton_kernel/kv_move.py b/lightllm/common/basemodel/triton_kernel/kv_move.py index ce99a007bd..34b812327b 100644 --- a/lightllm/common/basemodel/triton_kernel/kv_move.py +++ b/lightllm/common/basemodel/triton_kernel/kv_move.py @@ -42,6 +42,64 @@ def _copy_kv_buffer_to_kv_buffer( return +@triton.jit +def _copy_kv_buffer_pair_to_kv_buffer_pair( + mem_num, + src_mem_index, + dst_mem_index, + src_index_stride: tl.constexpr, + dst_index_stride: tl.constexpr, + k_buffer, + v_buffer, + k_buffer_stride_l, + k_buffer_stride_s, + k_buffer_stride_d, + v_buffer_stride_l, + v_buffer_stride_s, + v_buffer_stride_d, + k_buffer_tail_dim, + v_buffer_tail_dim, + BLOCK: tl.constexpr, +): + layer_index = tl.program_id(0).to(tl.int64) + start_index = tl.program_id(1).to(tl.int64) + grid_num = tl.num_programs(1).to(tl.int64) + + k_buffer_stride_l = tl.cast(k_buffer_stride_l, dtype=tl.int64) + k_buffer_stride_s = tl.cast(k_buffer_stride_s, dtype=tl.int64) + k_buffer_stride_d = tl.cast(k_buffer_stride_d, dtype=tl.int64) + v_buffer_stride_l = tl.cast(v_buffer_stride_l, dtype=tl.int64) + v_buffer_stride_s = tl.cast(v_buffer_stride_s, dtype=tl.int64) + v_buffer_stride_d = tl.cast(v_buffer_stride_d, dtype=tl.int64) + + for i in range(start_index, mem_num, grid_num): + src_mem = tl.load(src_mem_index + i * src_index_stride) + dst_mem = tl.load(dst_mem_index + i * dst_index_stride) + for j in range(tl.cdiv(k_buffer_tail_dim, BLOCK)): + offs = j * BLOCK + tl.arange(0, BLOCK) + mask = offs < k_buffer_tail_dim + k_buffer_data = tl.load( + k_buffer + layer_index * k_buffer_stride_l + src_mem * k_buffer_stride_s + offs, mask=mask + ) + tl.store( + k_buffer + layer_index * k_buffer_stride_l + dst_mem * k_buffer_stride_s + offs, + k_buffer_data, + mask=mask, + ) + for j in range(tl.cdiv(v_buffer_tail_dim, BLOCK)): + offs = j * BLOCK + tl.arange(0, BLOCK) + mask = offs < v_buffer_tail_dim + v_buffer_data = tl.load( + v_buffer + layer_index * v_buffer_stride_l + src_mem * v_buffer_stride_s + offs, mask=mask + ) + tl.store( + v_buffer + layer_index * v_buffer_stride_l + dst_mem * v_buffer_stride_s + offs, + v_buffer_data, + mask=mask, + ) + return + + def copy_kv_buffer_to_kv_buffer( src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor, @@ -70,3 +128,42 @@ def copy_kv_buffer_to_kv_buffer( BLOCK=BLOCK, ) return + + +def copy_kv_buffer_pair_to_kv_buffer_pair( + src_mem_index: torch.Tensor, + dst_mem_index: torch.Tensor, + k_buffer: torch.Tensor, + v_buffer: torch.Tensor, +): + assert len(src_mem_index) == len(dst_mem_index) + assert src_mem_index.is_cuda and dst_mem_index.is_cuda and k_buffer.is_cuda and v_buffer.is_cuda + assert k_buffer.shape == v_buffer.shape + k_buffer = k_buffer.view(k_buffer.shape[0], k_buffer.shape[1], -1).view(dtype=torch.uint8) + v_buffer = v_buffer.view(v_buffer.shape[0], v_buffer.shape[1], -1).view(dtype=torch.uint8) + assert k_buffer.shape == v_buffer.shape + BLOCK = 4096 + layer_num = k_buffer.shape[0] + grid = ( + layer_num, + 1024, + ) + _copy_kv_buffer_pair_to_kv_buffer_pair[grid]( + mem_num=len(src_mem_index), + src_mem_index=src_mem_index, + dst_mem_index=dst_mem_index, + src_index_stride=src_mem_index.stride(0), + dst_index_stride=dst_mem_index.stride(0), + k_buffer=k_buffer, + v_buffer=v_buffer, + k_buffer_stride_l=k_buffer.stride(0), + k_buffer_stride_s=k_buffer.stride(1), + k_buffer_stride_d=k_buffer.stride(2), + v_buffer_stride_l=v_buffer.stride(0), + v_buffer_stride_s=v_buffer.stride(1), + v_buffer_stride_d=v_buffer.stride(2), + k_buffer_tail_dim=k_buffer.shape[-1], + v_buffer_tail_dim=v_buffer.shape[-1], + BLOCK=BLOCK, + ) + return diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py index 2bf325340f..0055ea6d93 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py @@ -5,6 +5,38 @@ import torch +def _load_causal_conv1d_fwd(): + try: + from sgl_kernel import causal_conv1d_fwd + + return causal_conv1d_fwd + except (ImportError, OSError): + try: + from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( + causal_conv1d_fn as vllm_causal_conv1d_fn, + ) + + return vllm_causal_conv1d_fn + except (ImportError, OSError): + return None + + +def _load_causal_conv1d_update(): + try: + from sgl_kernel import causal_conv1d_update + + return causal_conv1d_update + except (ImportError, OSError): + try: + from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( + causal_conv1d_update, + ) + + return causal_conv1d_update + except (ImportError, OSError): + return None + + def causal_conv1d_fn( x: torch.Tensor, weight: torch.Tensor, @@ -17,54 +49,70 @@ def causal_conv1d_fn( pad_slot_id: int = -1, **kwargs, ): - """ - x: (batch, dim, seqlen) or (dim,cu_seq_len) for varlen - sequences are concatenated from left to right for varlen - weight: (dim, width) - bias: (dim,) - query_start_loc: (batch + 1) int32 - The cumulative sequence lengths of the sequences in - the batch, used to index into sequence. prepended by 0. - for example: query_start_loc = torch.Tensor([0,10,16,17]), - x.shape=(dim,17) - cache_indices: (batch) int32 - indicates the corresponding state index, - like so: conv_state = conv_states[cache_indices[batch_id]] - has_initial_state: (batch) bool - indicates whether should the kernel take the current state as initial - state for the calculations - conv_states: (...,dim,width - 1) itype - updated inplace if provided - activation: either None or "silu" or "swish" - pad_slot_id: int - if cache_indices is passed, lets the kernel identify padded - entries that will not be processed, - for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id] - in this case, the kernel will not process entries at - indices 0 and 3 - - - out: (batch, dim, seqlen) + """Run causal depthwise convolution with SGL Kernel or a local Triton fallback. + + x is [batch, channels, tokens] or packed [channels, total_tokens]. The Triton + path requires packed inputs and all four request/state tensors below. + weight is [channels, width]; optional bias is [channels]. + + query_start_loc: cumulative token offsets [batch + 1], starting at zero. + cache_indices: maps each request to a slot in conv_states; pad_slot_id + entries keep their input and state unchanged. + has_initial_state: one bool per request; false means a zero-filled prefix. + conv_states: [slots, channels, width - 1], updated in place with input tails. + + activation is None, "silu" or "swish". The optional max_seqlen keyword bounds + the longest packed query for Triton without reading GPU lengths on the CPU. + Returns the same shape as x. SGL may overwrite x; Triton allocates an output. """ if activation not in [None, "silu", "swish"]: raise NotImplementedError("activation must be None, silu, or swish") - from sgl_kernel import causal_conv1d_fwd + optimized_causal_conv1d = _load_causal_conv1d_fwd() + if optimized_causal_conv1d is None: + from .causal_conv1d_triton import causal_conv1d_fn as triton_causal_conv1d_fn + + return triton_causal_conv1d_fn( + x, + weight, + bias, + query_start_loc=query_start_loc, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + conv_states=conv_states, + activation=activation, + pad_slot_id=pad_slot_id, + **kwargs, + ) if x.stride(-1) != 1: x = x.contiguous() bias = bias.contiguous() if bias is not None else None - causal_conv1d_fwd( - x, - weight, - bias, - conv_states, - query_start_loc, - cache_indices, - has_initial_state, - activation in ["silu", "swish"], - pad_slot_id, - ) + if optimized_causal_conv1d.__module__.startswith("sgl_kernel"): + optimized_causal_conv1d( + x, + weight, + bias, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + activation in ["silu", "swish"], + pad_slot_id, + ) + else: + return optimized_causal_conv1d( + x, + weight, + bias, + conv_states, + query_start_loc, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + activation=activation, + pad_slot_id=pad_slot_id, + null_block_id=None, + ) return x @@ -78,46 +126,74 @@ def causal_conv1d_update( conv_state_indices: Optional[torch.Tensor] = None, pad_slot_id: int = -1, ): - """ - x: (batch, dim) or (batch, dim, seqlen) - conv_state: (batch, dim, state_len), where state_len >= width - 1 - weight: (dim, width) - bias: (dim,) - cache_seqlens: (batch,), dtype int32. - If not None, the conv_state is treated as a circular buffer. - The conv_state will be updated by copying x to the conv_state - starting at the index - @cache_seqlens % state_len. - conv_state_indices: (batch,), dtype int32 - If not None, the conv_state is a larger tensor along the batch dim, - and we are selecting the batch coords specified by conv_state_indices. - Useful for a continuous batching scenario. - pad_slot_id: int - if cache_indices is passed, lets the kernel identify padded - entries that will not be processed, - for example: cache_indices = [pad_slot_id, 1 ,20 ,pad_slot_id] - in this case, the kernel will not process entries at - indices 0 and 3 - out: (batch, dim) or (batch, dim, seqlen) + """Decode with SGL Kernel when available, otherwise the local Triton kernel. + + x is [batch, channels] or [batch, channels, tokens]. conv_state is + [slots, channels, state_len], where state_len >= width - 1, and is updated + in place. weight is [channels, width]; optional bias is [channels]. + + cache_seqlens supplies one cursor per request for circular state buffers; + new tokens are written starting at cursor % state_len. Without cursors, + states shift left and append the new tokens. + conv_state_indices optionally maps requests to cache slots. Entries equal + to pad_slot_id keep their input and state unchanged. + + activation is None, "silu" or "swish". Returns the same shape as x. + SGL updates x in place; Triton allocates an output. """ if activation not in [None, "silu", "swish"]: raise NotImplementedError(f"activation must be None, silu, or swish, actual: {activation}") - from sgl_kernel import causal_conv1d_update as causal_conv1d_update_kernel + optimized_causal_conv1d_update = _load_causal_conv1d_update() + if optimized_causal_conv1d_update is None: + from .causal_conv1d_triton import causal_conv1d_update as triton_causal_conv1d_update + + return triton_causal_conv1d_update( + x, + conv_state, + weight, + bias, + activation=activation, + cache_seqlens=cache_seqlens, + conv_state_indices=conv_state_indices, + pad_slot_id=pad_slot_id, + ) + + if cache_seqlens is not None: + from .causal_conv1d_triton import causal_conv1d_update as triton_causal_conv1d_update + + return triton_causal_conv1d_update( + x, + conv_state, + weight, + bias, + activation=activation, + cache_seqlens=cache_seqlens, + conv_state_indices=conv_state_indices, + pad_slot_id=pad_slot_id, + ) activation_val = activation in ["silu", "swish"] - unsqueeze = x.dim() == 2 - if unsqueeze: - x = x.unsqueeze(-1) - causal_conv1d_update_kernel( + if optimized_causal_conv1d_update.__module__.startswith("sgl_kernel"): + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + optimized_causal_conv1d_update( + x, + conv_state, + weight, + bias, + activation_val, + None, + conv_state_indices, + pad_slot_id=pad_slot_id, + ) + return x.squeeze(-1) if unsqueeze else x + return optimized_causal_conv1d_update( x, conv_state, weight, bias, - activation_val, - cache_seqlens, - conv_state_indices, - pad_slot_id, + activation=activation, + conv_state_indices=conv_state_indices, + null_block_id=pad_slot_id, ) - if unsqueeze: - x = x.squeeze(-1) - return x diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_triton.py b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_triton.py new file mode 100644 index 0000000000..162cb36c60 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_triton.py @@ -0,0 +1,435 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Causal depthwise convolution without external CUDA extensions. + +Packed prefill reads strided [channels, total_tokens] projection views. +Outputs are allocated separately; convolution states are updated in place. +""" + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune + + +def _prefill_configs(): + return [ + {"BLOCK_TOKENS": block_tokens, "BLOCK_CHANNELS": block_channels, "num_warps": num_warps} + for block_tokens in (8, 16, 32, 64) + for block_channels in (64, 128, 256) + for num_warps in (2, 4, 8) + if block_tokens * block_channels <= num_warps * 32 * 64 + ] + + +def _update_configs(): + return [ + {"BLOCK_CHANNELS": block_channels, "num_warps": num_warps} + for block_channels in (64, 128, 256) + for num_warps in (1, 2, 4) + ] + + +def _conv_static_key(x, weight, bias, activation, conv_state): + # Group by memory order so nearby token counts can reuse a configuration. + if x.stride(-2) == 1: + layout = "channel" + elif x.stride(-1) == 1: + layout = "token" + else: + layout = "strided" + return { + "channels": weight.shape[0], + "width": weight.shape[1], + "dtype": str(x.dtype), + "layout": layout, + "state_stride": f"{conv_state.stride(-2)}_{conv_state.stride(-1)}", + "bias": bias is not None, + "silu": activation is not None, + } + + +def _prefill_static_key(x, weight, bias, conv_states, activation): + return _conv_static_key(x, weight, bias, activation, conv_states) + + +def _prefill_run_key(x, query_start_loc, max_seqlen): + batch_size = query_start_loc.numel() - 1 + max_query_len = x.shape[-1] if max_seqlen is None else max_seqlen + return f"{batch_size}_{max_query_len}" + + +def _prefill_key_distance(run_key, config_key): + batch_size, max_query_len = map(int, run_key.split("_")) + config_batch_size, config_max_query_len = map(int, config_key.split("_")) + return ( + abs(batch_size * max_query_len - config_batch_size * config_max_query_len), + abs(batch_size - config_batch_size), + ) + + +def _update_static_key(x, conv_state, weight, bias, activation, cache_seqlens): + input_3d = x.unsqueeze(-1) if x.ndim == 2 else x + key = _conv_static_key(input_3d, weight, bias, activation, conv_state) + key.update( + tokens=input_3d.shape[-1], + state_len=conv_state.shape[-1], + circular=cache_seqlens is not None, + ) + return key + + +@triton.jit +def _causal_conv1d_fwd_kernel( + x_ptr, + weight_ptr, + bias_ptr, + conv_states_ptr, + query_start_loc_ptr, + cache_indices_ptr, + has_initial_state_ptr, + out_ptr, + num_channels, + stride_x_channel, + stride_x_token, + stride_weight_channel, + stride_weight_tap, + stride_state_cache, + stride_state_channel, + stride_state_token, + stride_out_channel, + stride_out_token, + CONV_WIDTH: tl.constexpr, + HAS_BIAS: tl.constexpr, + APPLY_SILU: tl.constexpr, + PAD_SLOT_ID: tl.constexpr, + BLOCK_TOKENS: tl.constexpr, + BLOCK_CHANNELS: tl.constexpr, +): + HISTORY_LEN: tl.constexpr = CONV_WIDTH - 1 + tl.static_assert(CONV_WIDTH >= 2 and CONV_WIDTH <= 4, "Supported convolution widths are 2, 3 and 4") + tl.static_assert(BLOCK_TOKENS >= HISTORY_LEN, "Only the first token block may read initial state") + + sequence_index = tl.program_id(0) + token_block_index = tl.program_id(1) + channel_indices = tl.program_id(2) * BLOCK_CHANNELS + tl.arange(0, BLOCK_CHANNELS) + valid_channels = channel_indices < num_channels + sequence_start = tl.load(query_start_loc_ptr + sequence_index) + sequence_length = tl.load(query_start_loc_ptr + sequence_index + 1) - sequence_start + if token_block_index * BLOCK_TOKENS >= sequence_length: + return + + token_indices = token_block_index * BLOCK_TOKENS + tl.arange(0, BLOCK_TOKENS) + valid_output = (token_indices[:, None] < sequence_length) & valid_channels[None, :] + input_base = x_ptr + sequence_start * stride_x_token + channel_indices * stride_x_channel + output_base = out_ptr + sequence_start * stride_out_token + channel_indices * stride_out_channel + output_ptrs = output_base[None, :] + token_indices[:, None] * stride_out_token + cache_index = tl.load(cache_indices_ptr + sequence_index) + if cache_index == PAD_SLOT_ID: + # Padding requests retain their input, as in the CUDA implementation. + padded_input = tl.load(input_base[None, :] + token_indices[:, None] * stride_x_token, valid_output, other=0) + tl.store(output_ptrs, padded_input, valid_output) + return + + state_base = conv_states_ptr + cache_index * stride_state_cache + channel_indices * stride_state_channel + weight_base = weight_ptr + channel_indices * stride_weight_channel + use_initial_state = tl.load(has_initial_state_ptr + sequence_index) + accumulator = tl.full((BLOCK_TOKENS, BLOCK_CHANNELS), 0, tl.float32) + if HAS_BIAS: + bias_values = tl.load(bias_ptr + channel_indices, valid_channels, other=0).to(tl.float32) + accumulator += bias_values[None, :] + + # Negative token positions come from cached history, or zero for a new request. + for tap in tl.static_range(CONV_WIDTH): + source_tokens = token_indices[:, None] + tap - HISTORY_LEN + input_values = tl.load( + input_base[None, :] + source_tokens * stride_x_token, valid_output & (source_tokens >= 0), other=0 + ).to(tl.float32) + history_values = tl.load( + state_base[None, :] + (source_tokens + HISTORY_LEN) * stride_state_token, + valid_output & (source_tokens < 0) & use_initial_state, + other=0, + ).to(tl.float32) + window_values = tl.where(source_tokens < 0, history_values, input_values) + conv_weight = tl.load(weight_base + tap * stride_weight_tap, valid_channels, other=0).to(tl.float32) + accumulator = tl.fma(window_values, conv_weight[None, :], accumulator) + if APPLY_SILU: + accumulator = accumulator / (1.0 + tl.exp(-accumulator)) + tl.store(output_ptrs, accumulator, valid_output) + + if token_block_index == 0: + # Only the first token block reads/writes initial state. Later blocks + # obtain their history from input, avoiding cross-CTA state races. + state_offsets = tl.arange(0, triton.next_power_of_2(HISTORY_LEN))[:, None] + tail_tokens = sequence_length - HISTORY_LEN + state_offsets + valid_state = (state_offsets < HISTORY_LEN) & valid_channels[None, :] + tail_inputs = tl.load( + input_base[None, :] + tail_tokens * stride_x_token, valid_state & (tail_tokens >= 0), other=0 + ) + tail_history = tl.load( + state_base[None, :] + (tail_tokens + HISTORY_LEN) * stride_state_token, + valid_state & (tail_tokens < 0) & use_initial_state, + other=0, + ) + updated_state = tl.where(tail_tokens >= 0, tail_inputs, tail_history) + # Finish all prefix reads before overwriting their state slots, + # including requests shorter than the history window. + tl.debug_barrier() + tl.store(state_base[None, :] + state_offsets * stride_state_token, updated_state, valid_state) + + +@autotune( + kernel_name="causal_conv1d_prefill:v1", + configs_gen_func=_prefill_configs, + static_key_func=_prefill_static_key, + run_key_func=_prefill_run_key, + run_key_distance_func=_prefill_key_distance, + mutates_args=["conv_states"], +) +def causal_conv1d_fn( + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + *, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor, + has_initial_state: torch.Tensor, + conv_states: torch.Tensor, + activation: Optional[str] = "silu", + pad_slot_id: int = -1, + max_seqlen: Optional[int] = None, + run_config: Optional[dict] = None, + **kwargs, +): + """Convolve packed [channels, total_tokens] inputs and update mapped states. + + Request boundaries, cache indices and initial-state flags are required, + matching the GDN prefill caller. States have shape [slots, channels, width - 1] + and may be strided MTP slices. Metadata tensors are contiguous. + max_seqlen bounds the longest query; total tokens are a safe default that + avoids reading lengths back from the GPU. Empty/padding requests keep state. + """ + if activation not in (None, "silu", "swish"): + raise NotImplementedError(f"Unsupported activation: {activation}") + num_channels, total_tokens = x.shape + batch_size = query_start_loc.numel() - 1 + out = torch.empty_like(x) + if batch_size == 0 or x.numel() == 0: + return out + + bias = bias.contiguous() if bias is not None else None + max_query_len = total_tokens if max_seqlen is None else max_seqlen + if run_config is None: + run_config = {"BLOCK_TOKENS": 16, "BLOCK_CHANNELS": 128, "num_warps": 4} + block_tokens = run_config["BLOCK_TOKENS"] + block_channels = run_config["BLOCK_CHANNELS"] + grid = (batch_size, triton.cdiv(max_query_len, block_tokens), triton.cdiv(num_channels, block_channels)) + _causal_conv1d_fwd_kernel[grid]( + x, + weight, + bias, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + out, + num_channels, + *x.stride(), + *weight.stride(), + *conv_states.stride(), + *out.stride(), + CONV_WIDTH=weight.shape[1], + HAS_BIAS=bias is not None, + APPLY_SILU=activation is not None, + PAD_SLOT_ID=pad_slot_id, + BLOCK_TOKENS=block_tokens, + BLOCK_CHANNELS=block_channels, + num_warps=run_config["num_warps"], + ) + return out + + +@triton.jit +def _causal_conv1d_update_kernel( + x_ptr, + weight_ptr, + bias_ptr, + conv_state_ptr, + cache_seqlens_ptr, + conv_state_indices_ptr, + out_ptr, + num_channels, + stride_x_batch, + stride_x_channel, + stride_x_token, + stride_weight_channel, + stride_weight_tap, + stride_state_cache, + stride_state_channel, + stride_state_token, + stride_out_batch, + stride_out_channel, + stride_out_token, + CONV_WIDTH: tl.constexpr, + NUM_TOKENS: tl.constexpr, + STATE_LEN: tl.constexpr, + HAS_INDICES: tl.constexpr, + HAS_BIAS: tl.constexpr, + CIRCULAR_BUFFER: tl.constexpr, + APPLY_SILU: tl.constexpr, + PAD_SLOT_ID: tl.constexpr, + BLOCK_CHANNELS: tl.constexpr, +): + tl.static_assert(CONV_WIDTH >= 2 and CONV_WIDTH <= 4, "Supported convolution widths are 2, 3 and 4") + tl.static_assert(STATE_LEN >= CONV_WIDTH - 1, "State must hold the full convolution history") + + batch_index = tl.program_id(0) + channel_indices = tl.program_id(1) * BLOCK_CHANNELS + tl.arange(0, BLOCK_CHANNELS) + valid_channels = channel_indices < num_channels + input_base = x_ptr + batch_index * stride_x_batch + channel_indices * stride_x_channel + output_base = out_ptr + batch_index * stride_out_batch + channel_indices * stride_out_channel + cache_index = tl.load(conv_state_indices_ptr + batch_index) if HAS_INDICES else batch_index + if cache_index == PAD_SLOT_ID: + for token_index in range(NUM_TOKENS): + padded_input = tl.load(input_base + token_index * stride_x_token, valid_channels, other=0) + tl.store(output_base + token_index * stride_out_token, padded_input, valid_channels) + return + state_base = conv_state_ptr + cache_index * stride_state_cache + channel_indices * stride_state_channel + weight_base = weight_ptr + channel_indices * stride_weight_channel + state_cursor = tl.load(cache_seqlens_ptr + batch_index) % STATE_LEN if CIRCULAR_BUFFER else STATE_LEN + history_start = state_cursor + STATE_LEN - (CONV_WIDTH - 1) + + # Keep the short convolution window and weights in registers across tokens. + history_0 = tl.load( + state_base + (history_start % STATE_LEN) * stride_state_token, + valid_channels, + other=0, + ).to(tl.float32) + weight_0 = tl.load(weight_base, valid_channels, other=0).to(tl.float32) + weight_1 = tl.load(weight_base + stride_weight_tap, valid_channels, other=0).to(tl.float32) + if CONV_WIDTH >= 3: + history_1 = tl.load( + state_base + ((history_start + 1) % STATE_LEN) * stride_state_token, + valid_channels, + other=0, + ).to(tl.float32) + weight_2 = tl.load(weight_base + 2 * stride_weight_tap, valid_channels, other=0).to(tl.float32) + if CONV_WIDTH == 4: + history_2 = tl.load( + state_base + ((history_start + 2) % STATE_LEN) * stride_state_token, valid_channels, other=0 + ).to(tl.float32) + weight_3 = tl.load(weight_base + 3 * stride_weight_tap, valid_channels, other=0).to(tl.float32) + bias_values = tl.full((BLOCK_CHANNELS,), 0, tl.float32) + if HAS_BIAS: + bias_values = tl.load(bias_ptr + channel_indices, valid_channels, other=0).to(tl.float32) + + for token_index in range(NUM_TOKENS): + input_values = tl.load(input_base + token_index * stride_x_token, valid_channels, other=0).to(tl.float32) + accumulator = tl.fma(weight_0, history_0, bias_values) + if CONV_WIDTH == 2: + accumulator = tl.fma(weight_1, input_values, accumulator) + history_0 = input_values + elif CONV_WIDTH == 3: + accumulator = tl.fma(weight_1, history_1, accumulator) + accumulator = tl.fma(weight_2, input_values, accumulator) + history_0, history_1 = history_1, input_values + else: + accumulator = tl.fma(weight_1, history_1, accumulator) + accumulator = tl.fma(weight_2, history_2, accumulator) + accumulator = tl.fma(weight_3, input_values, accumulator) + history_0, history_1, history_2 = history_1, history_2, input_values + if APPLY_SILU: + accumulator = accumulator / (1.0 + tl.exp(-accumulator)) + tl.store(output_base + token_index * stride_out_token, accumulator, valid_channels) + if CIRCULAR_BUFFER: + tl.store( + state_base + ((state_cursor + token_index) % STATE_LEN) * stride_state_token, + input_values, + valid_channels, + ) + if not CIRCULAR_BUFFER: + # Keep the last STATE_LEN tokens of [old state, new input]. + state_offsets = tl.arange(0, triton.next_power_of_2(STATE_LEN))[:, None] + valid_state = (state_offsets < STATE_LEN) & valid_channels[None, :] + source_offsets = state_offsets + NUM_TOKENS + old_history = tl.load( + state_base[None, :] + source_offsets * stride_state_token, + valid_state & (source_offsets < STATE_LEN), + other=0, + ) + new_inputs = tl.load( + input_base[None, :] + (source_offsets - STATE_LEN) * stride_x_token, + valid_state & (source_offsets >= STATE_LEN), + other=0, + ) + updated_state = tl.where(source_offsets < STATE_LEN, old_history, new_inputs) + tl.debug_barrier() + tl.store(state_base[None, :] + state_offsets * stride_state_token, updated_state, valid_state) + + +@autotune( + kernel_name="causal_conv1d_update:v1", + kernel_type=AutotuneKernelType.DECODE_ATTENTION, + configs_gen_func=_update_configs, + static_key_func=_update_static_key, + run_key_func=lambda x: x.shape[0], + mutates_args=["conv_state"], +) +def causal_conv1d_update( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + activation: Optional[str] = None, + cache_seqlens: Optional[torch.Tensor] = None, + conv_state_indices: Optional[torch.Tensor] = None, + pad_slot_id: int = -1, + run_config: Optional[dict] = None, +): + """Decode [batch, channels] or [batch, channels, tokens], updating linear/ring states. + + State shape is [slots, channels, state_len], with state_len >= width - 1. + Cache indices and sequence lengths, when supplied, are contiguous batch vectors. + """ + if activation not in (None, "silu", "swish"): + raise NotImplementedError(f"Unsupported activation: {activation}") + out = torch.empty_like(x) + if x.numel() == 0: + return out + + input_3d = x.unsqueeze(-1) if x.ndim == 2 else x + output_3d = out.unsqueeze(-1) if x.ndim == 2 else out + batch_size, num_channels, num_tokens = input_3d.shape + bias = bias.contiguous() if bias is not None else None + if run_config is None: + run_config = {"BLOCK_CHANNELS": 128, "num_warps": 4} + block_channels = run_config["BLOCK_CHANNELS"] + grid = (batch_size, triton.cdiv(num_channels, block_channels)) + _causal_conv1d_update_kernel[grid]( + input_3d, + weight, + bias, + conv_state, + cache_seqlens, + conv_state_indices, + output_3d, + num_channels, + *input_3d.stride(), + *weight.stride(), + *conv_state.stride(), + *output_3d.stride(), + CONV_WIDTH=weight.shape[1], + NUM_TOKENS=num_tokens, + STATE_LEN=conv_state.shape[2], + HAS_INDICES=conv_state_indices is not None, + HAS_BIAS=bias is not None, + CIRCULAR_BUFFER=cache_seqlens is not None, + APPLY_SILU=activation is not None, + PAD_SLOT_ID=pad_slot_id, + BLOCK_CHANNELS=block_channels, + num_warps=run_config["num_warps"], + ) + return out diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py index cd3b0962a3..1fd7fabe3f 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py @@ -7,9 +7,10 @@ # the following copyright notice: # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang from .chunk import chunk_gated_delta_rule -from .fused_recurrent import fused_recurrent_gated_delta_rule +from .fused_recurrent import fused_recurrent_gated_delta_rule, fused_recurrent_gated_delta_rule_packed_decode __all__ = [ "chunk_gated_delta_rule", "fused_recurrent_gated_delta_rule", + "fused_recurrent_gated_delta_rule_packed_decode", ] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py index 5dfbd6e4ab..53b29d0437 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py @@ -380,6 +380,171 @@ def fused_recurrent_gated_delta_rule_fwd( return o, final_state +@triton.jit +def _fused_recurrent_gated_delta_rule_packed_decode_kernel( + mixed_qkv, + a, + b, + A_log, + dt_bias, + o, + h0, + ssm_state_indices, + scale, + stride_mixed_qkv_tok: tl.constexpr, + stride_a_tok: tl.constexpr, + stride_b_tok: tl.constexpr, + stride_state_slot: tl.constexpr, + stride_indices: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + SOFTPLUS_THRESHOLD: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, +): + """Packed one-token GDN decode for the non-MTP path. + + ``mixed_qkv`` stays in its post-convolution ``[q, k, v]`` layout. Each + program owns one value-head tile, so the kernel avoids materializing + reshaped q/k/v tensors and writes the recurrent state in-place. LightLLM + stores the recurrent state as ``[slot, head, key, value]``; this differs + from the ``[slot, head, value, key]`` layout used by the upstream FLA + kernel, so the state offsets intentionally keep key as the first matrix + dimension. + """ + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + state_idx = tl.load(ssm_state_indices + i_n * stride_indices).to(tl.int64) + p_h0 = h0 + state_idx * stride_state_slot + p_h0 += i_hv * K * V + o_k[:, None] * V + o_v[None, :] + b_h = tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + p_mixed = mixed_qkv + i_n * stride_mixed_qkv_tok + q_off = i_h * K + o_k + k_off = H * K + i_h * K + o_k + v_off = 2 * H * K + i_hv * V + o_v + b_q = tl.load(p_mixed + q_off, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_mixed + k_off, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_mixed + v_off, mask=mask_v, other=0).to(tl.float32) + + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q = b_q * scale + + a_val = tl.load(a + i_n * stride_a_tok + i_hv).to(tl.float32) + b_val = tl.load(b + i_n * stride_b_tok + i_hv).to(tl.float32) + A_log_val = tl.load(A_log + i_hv).to(tl.float32) + dt_bias_val = tl.load(dt_bias + i_hv).to(tl.float32) + x = a_val + dt_bias_val + softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x) + g_val = -tl.exp(A_log_val) * softplus_x + beta_val = tl.sigmoid(b_val) + + b_h *= exp(g_val) + b_v -= tl.sum(b_h * b_k[:, None], 0) + b_v *= beta_val + b_h += b_k[:, None] * b_v[None, :] + b_o = tl.sum(b_h * b_q[:, None], 0) + + p_o = o + (i_n * HV + i_hv) * V + o_v + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h) + + +def fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + ssm_state_indices: torch.Tensor, + use_qk_l2norm_in_kernel: bool = False, + out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run packed one-token GDN decode without q/k/v materialization.""" + if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1: + raise ValueError("mixed_qkv must be a 2D tensor contiguous in its last dimension") + if a.ndim != 2 or b.ndim != 2 or a.stride(-1) != 1 or b.stride(-1) != 1: + raise ValueError("a and b must be 2D tensors contiguous in their last dimension") + if A_log.ndim != 1 or dt_bias.ndim != 1 or A_log.stride(0) != 1 or dt_bias.stride(0) != 1: + raise ValueError("A_log and dt_bias must be contiguous 1D tensors") + if ssm_state_indices.ndim != 1 or not ssm_state_indices.is_contiguous(): + raise ValueError("ssm_state_indices must be a contiguous 1D tensor") + if initial_state.ndim != 4 or initial_state.stride(-1) != 1: + raise ValueError("initial_state must be a 4D tensor contiguous in its last dimension") + + batch = mixed_qkv.shape[0] + num_v_heads, key_dim, value_dim = initial_state.shape[-3:] + if a.shape != (batch, num_v_heads) or b.shape != (batch, num_v_heads): + raise ValueError("a and b must have shape [batch, num_v_heads]") + if ssm_state_indices.shape[0] != batch: + raise ValueError("ssm_state_indices must have one entry per batch row") + if A_log.numel() != num_v_heads or dt_bias.numel() != num_v_heads: + raise ValueError("A_log and dt_bias must match num_v_heads") + + qk_dim = mixed_qkv.shape[1] - num_v_heads * value_dim + if qk_dim <= 0 or qk_dim % 2 != 0: + raise ValueError("mixed_qkv does not contain a valid q/k/v layout") + query_dim = qk_dim // 2 + if query_dim % key_dim != 0: + raise ValueError("query dimension must be divisible by key head dimension") + num_k_heads = query_dim // key_dim + if num_k_heads <= 0 or num_v_heads % num_k_heads != 0: + raise ValueError("invalid GQA head configuration") + + block_k = triton.next_power_of_2(key_dim) + if triton.cdiv(key_dim, block_k) != 1: + raise ValueError("packed decode currently supports one key tile") + block_v = min(triton.next_power_of_2(value_dim), 32) + if out is None: + out = mixed_qkv.new_empty((batch, num_v_heads, value_dim)) + elif out.shape != (batch, num_v_heads, value_dim) or not out.is_contiguous(): + raise ValueError("out must be contiguous with shape [batch, num_v_heads, value_dim]") + + grid = (triton.cdiv(value_dim, block_v), batch * num_v_heads) + _fused_recurrent_gated_delta_rule_packed_decode_kernel[grid]( + mixed_qkv=mixed_qkv, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + o=out, + h0=initial_state, + ssm_state_indices=ssm_state_indices, + scale=scale, + stride_mixed_qkv_tok=mixed_qkv.stride(0), + stride_a_tok=a.stride(0), + stride_b_tok=b.stride(0), + stride_state_slot=initial_state.stride(0), + stride_indices=ssm_state_indices.stride(0), + H=num_k_heads, + HV=num_v_heads, + K=key_dim, + V=value_dim, + BK=block_k, + BV=block_v, + SOFTPLUS_THRESHOLD=20.0, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + num_warps=1, + num_stages=3, + ) + return out.unsqueeze(1), initial_state + + class FusedRecurrentFunction(torch.autograd.Function): @staticmethod def forward( diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py b/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py index b81bf9ec16..0413aea898 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py @@ -3,7 +3,7 @@ import triton.language as tl -@triton.jit +@triton.jit(do_not_specialize=["batch_size"]) def _build_dynamic_mtp_linear_att_state_params_kernel( b_req_idx, b_mtp_index, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py new file mode 100644 index 0000000000..e54186eae4 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +"""ReplaySSM kernel exports, in lifecycle order. + +decode writes history for the current window. +cursor records how many of those tokens were accepted. +fold applies that accepted prefix onto the canonical SSM. +""" + +from .replayssm_decode import ( + _decode_replay, + _gdn_precomputed_window, + _kda_precomputed_window, + _project_qk, + _replay_each_token, +) +from .replayssm_delta import _delta_rule_step +from .replayssm_cursor import _advance_cursor +from .replayssm_fold import _fold_prefix + +__all__ = [ + "_decode_replay", + "_project_qk", + "_gdn_precomputed_window", + "_kda_precomputed_window", + "_replay_each_token", + "_delta_rule_step", + "_advance_cursor", + "_fold_prefix", +] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py new file mode 100644 index 0000000000..5354c71260 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Advance the accepted-prefix cursor. This kernel does not fold state.""" + +import triton +import triton.language as tl + + +@triton.jit(do_not_specialize=["N"]) +def _advance_cursor( + cursors, + reqs, + accepted_state_indices, + N, + HOLD: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK: tl.constexpr, +): + # N 是运行时请求数；BLOCK/WIDTH 仍是编译期参数，避免 eager batch 变化触发重新编译。 + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + r = tl.load(reqs + i, i < N, HOLD) + valid = (i < N) & (r != HOLD) + cursor = tl.load(cursors + r, valid, 0) + count = tl.load(accepted_state_indices + r, valid, 0) + 1 if WIDTH > 1 else 1 + tl.store(cursors + r, cursor + count, valid) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py new file mode 100644 index 0000000000..5619869b64 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py @@ -0,0 +1,792 @@ +# SPDX-License-Identifier: Apache-2.0 +"""ReplaySSM decode kernels. + +`_decode_replay` is the entry. It projects the verify window when asked, then +solves a GDN or KDA window, or walks tokens one by one. Neither path writes the +canonical SSM. +""" + +import triton +import triton.language as tl + + +@triton.jit +def _history_coeff(rows, vec, BK: tl.constexpr): + # rows [R, BK] · vec [BK] -> [R]. tf32x3 是 tensor core 上最接近 fp32 的输出重建。 + if BK >= 16: + return tl.sum(tl.dot(rows, vec[:, None], input_precision="tf32x3"), 1) + return tl.sum(rows * vec[None, :], 1) + + +@triton.jit +def _apply_deltas(coeff, deltas, R: tl.constexpr, BV: tl.constexpr): + # coeff [R] · deltas [R, BV] -> [BV]. R 小于 16 时 tensor core 不能做这次点积。 + if R >= 16 and BV >= 16: + return tl.sum(tl.dot(coeff[None, :], deltas, input_precision="tf32x3"), 0) + return tl.sum(deltas * coeff[:, None], 0) + + +@triton.jit +def _decode_replay( + Q, + Kp, + Vp, + A, + B, + Alog, + Bias, + State, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Reqs, + Cursors, + Cu, + Out, + SQ: tl.constexpr, + SK: tl.constexpr, + SV: tl.constexpr, + SA: tl.constexpr, + SB: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + HOLD: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + VARLEN: tl.constexpr, + PRECOMPUTE_PROJECTIONS: tl.constexpr, + BW: tl.constexpr, + KDA: tl.constexpr, + LOWER: tl.constexpr, + BQK: tl.constexpr, + BR: tl.constexpr, +): + """执行一个请求窗口的 replay decode，并写入 history。 + + 只把本步 token 追加到长度为 L 的连续前缀，不写 canonical SSM。 + cursor 是已有前缀长度，本步从那里接着写。请求数属于运行时布局， + 不应因为 eager batch 改变而形成新的 kernel 变体。 + """ + iv, seq, hv = tl.program_id(0), tl.program_id(1), tl.program_id(2) + req = tl.load(Reqs + seq).to(tl.int64) + if VARLEN: + start, end = tl.load(Cu + seq), tl.load(Cu + seq + 1) + else: + start, end = seq, seq + 1 + vv = iv * BV + tl.arange(0, BV) + if req == HOLD: + for t in range(start, end): + tl.store(Out + (t * HV + hv) * V + vv, 0, vv < V) + return + if start == end: + return + + kk = tl.arange(0, BK) + ll = tl.arange(0, L) + slot = req * HV + hv + key_slot = req * KH + hv // (HV // KH) + n = tl.load(Cursors + req) + sp = State + slot * K * V + kk[:, None] * V + vv[None, :] + state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) + + if PRECOMPUTE_PROJECTIONS: + # 输出用 tf32x3 tensor core 从冻结的 checkpoint 和 history 重建。 + # 状态不在这里更新；压缩仍用原始 k/v 做 fp32 串行折回。 + projection_dtype: tl.constexpr = tl.float32 + projection_precision: tl.constexpr = "tf32x3" + qk, projections, tt, tokens, valid, h = _project_qk( + Q, + Kp, + A, + Alog, + Bias, + Gates, + SQ, + SK, + SA, + H, + HV, + K, + L, + BW, + KDA, + LOWER, + BQK, + end, + hv, + kk, + ll, + n, + projection_dtype, + projection_precision, + slot, + start, + state, + ) + else: + projections, tt = None, None + + if PRECOMPUTE_PROJECTIONS and KDA: + _kda_precomputed_window( + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SK, + SV, + SA, + SB, + HV, + KH, + K, + V, + L, + BV, + BW, + BR, + BK, + LOWER, + end, + h, + hv, + iv, + key_slot, + kk, + n, + projections, + qk, + slot, + start, + tt, + vv, + ) + elif PRECOMPUTE_PROJECTIONS: + _gdn_precomputed_window( + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SK, + SV, + SA, + SB, + HV, + KH, + K, + V, + L, + BV, + BW, + BR, + end, + h, + hv, + iv, + key_slot, + kk, + n, + projection_dtype, + projection_precision, + projections, + qk, + slot, + start, + tokens, + tt, + valid, + vv, + ) + else: + _replay_each_token( + Q, + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SQ, + SK, + SV, + SA, + SB, + H, + HV, + KH, + K, + V, + L, + PRECOMPUTE_PROJECTIONS, + BW, + KDA, + LOWER, + BK, + BV, + end, + hv, + iv, + key_slot, + kk, + ll, + n, + projections, + slot, + start, + state, + tt, + vv, + ) + + +@triton.jit +def _project_qk( + Q, + Kp, + A, + Alog, + Bias, + Gates, + SQ: tl.constexpr, + SK: tl.constexpr, + SA: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + L: tl.constexpr, + BW: tl.constexpr, + KDA: tl.constexpr, + LOWER: tl.constexpr, + BQK: tl.constexpr, + end, + hv, + kk, + ll, + n, + projection_dtype: tl.constexpr, + projection_precision: tl.constexpr, + slot, + start, + state, +): + # 预计算 verify window 的 q/k 投影。 + tt = tl.arange(0, BQK) + tokens = start + tt % BW + valid = (tt < 2 * BW) & (tokens < end) + safe_tokens = tl.where(valid, tokens, start) + h = hv // (HV // H) + q_rows = tl.load( + Q + safe_tokens[:, None] * SQ + h * K + kk[None, :], + (valid & (tt < BW))[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + k_rows = tl.load( + Kp + safe_tokens[:, None] * SK + h * K + kk[None, :], + (valid & (tt >= BW))[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + # Padded query/key rows share the MMA tile. Keep them at zero so they + # cannot change the valid rows. + row_valid = valid[:, None] + qk = tl.where(row_valid, q_rows + k_rows, 0) + qk *= tl.rsqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] + qk *= tl.where(tt < BW, K ** -0.5, 1.0)[:, None] + qk = tl.where(row_valid, qk, 0) + if KDA: + history = ll < n + gates = tl.load( + Gates + (slot * L + ll[:, None]) * K + kk[None, :], + history[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + total_g = tl.sum(gates, axis=0) + wi = tl.arange(0, BW) + gate_ok = start + wi < end + safe_wi = tl.where(gate_ok, wi, 0) + gate_inputs = tl.load( + A + (start + safe_wi[:, None]) * SA + hv * K + kk[None, :], + gate_ok[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + gate_inputs += tl.load(Bias + hv * K + kk, kk < K, 0).to(tl.float32)[None, :] + fresh_gates = LOWER * tl.sigmoid(tl.exp(tl.load(Alog + hv)) * gate_inputs) + fresh_gates = tl.where((start + wi < end)[:, None], fresh_gates, 0) + query_gates = total_g[None, :] + tl.gather( + tl.cumsum(fresh_gates, axis=0), + tl.broadcast_to((tt % BW)[:, None], qk.shape), + 0, + ) + # All exponents are forward decay. No inverse prefix scaling can + # overflow when a long history contains strongly negative gates. + # Padded rows stay zero: 0 * exp(garbage) is NaN and would enter the dot. + scaled = tl.where(row_valid, qk * tl.exp(query_gates), 0) + projections = tl.dot( + scaled.to(projection_dtype), + state.to(projection_dtype), + input_precision=projection_precision, + ) + else: + projections = tl.dot(qk.to(projection_dtype), state.to(projection_dtype), input_precision=projection_precision) + return qk, projections, tt, tokens, valid, h + + +@triton.jit +def _gdn_precomputed_window( + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SK: tl.constexpr, + SV: tl.constexpr, + SA: tl.constexpr, + SB: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + BV: tl.constexpr, + BW: tl.constexpr, + BR: tl.constexpr, + end, + h, + hv, + iv, + key_slot, + kk, + n, + projection_dtype: tl.constexpr, + projection_precision: tl.constexpr, + projections, + qk, + slot, + start, + tokens, + tt, + valid, + vv, +): + # GDN precompute 路径：一次投影整个窗口，再按 token 顺序写入新 history。 + # 只有 corrected-value solve 依赖 token 顺序，避免重复计算历史投影。 + jj = tl.arange(0, BR) + fresh_tokens = start + jj - L + old = (jj < L) & (jj < n) + fresh = (jj >= L) & (jj < L + end - start) + # BR is rounded above the history length. Masked lanes must still address + # a real slot; otherwise they read uninitialized neighbor rows into the dot. + safe_jj = tl.where(jj < L, jj, 0) + safe_fresh = tl.where(fresh, fresh_tokens, start) + cached_keys = tl.load( + Keys + (key_slot * L) * K + safe_jj[:, None] * K + kk[None, :], + old[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + raw_keys = tl.load( + Kp + safe_fresh[:, None] * SK + h * K + kk[None, :], + fresh[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + fresh_keys = raw_keys * tl.rsqrt(tl.sum(raw_keys * raw_keys, axis=1) + 1.0e-6)[:, None] + fresh_keys = tl.where(fresh[:, None], fresh_keys, 0) + all_keys = cached_keys + fresh_keys + values = tl.load( + Vp + safe_fresh[:, None] * SV + hv * V + vv[None, :], + fresh[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + log_a, bias = tl.load(Alog + hv).to(tl.float32), tl.load(Bias + hv).to(tl.float32) + x = tl.load(A + safe_fresh * SA + hv, fresh, 0).to(tl.float32) + bias + fresh_gates = -tl.exp(log_a) * tl.where(x <= 20.0, tl.log(1.0 + tl.exp(x)), x) + gates = tl.load(Gates + slot * L + safe_jj, old, 0).to(tl.float32) + gates += tl.where(fresh, fresh_gates, 0.0) + prefix = tl.cumsum(gates, axis=0) + query_gates = tl.sum(tl.where(jj[None, :] <= L + (tt % BW)[:, None], gates[None, :], 0), 1) + coefficients = tl.dot( + qk.to(projection_dtype), tl.trans(all_keys).to(projection_dtype), input_precision=projection_precision + ) + causal = old[None, :] | (fresh[None, :] & (jj[None, :] <= L + (tt % BW)[:, None])) + coefficients *= tl.where(causal, tl.exp(tl.minimum(query_gates[:, None] - prefix[None, :], 0.0)), 0.0) + projections *= tl.exp(query_gates)[:, None] + deltas = tl.load( + Deltas + (slot * L) * V + safe_jj[:, None] * V + vv[None, :], + old[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + betas = tl.sigmoid(tl.load(B + safe_fresh * SB + hv, fresh, 0).to(tl.float32)) + for t in range(end - start): + checkpoint_sk = tl.sum(tl.where((tt == BW + t)[:, None], projections, 0), 0) + weights = tl.sum(tl.where((tt == BW + t)[:, None], coefficients, 0), 0) + sk = checkpoint_sk + _apply_deltas(weights, deltas, BR, BV) + value = tl.sum(tl.where((jj == L + t)[:, None], values, 0), 0) + beta = tl.sum(tl.where(jj == L + t, betas, 0)) + delta = beta * (value - sk) + deltas = tl.where((jj == L + t)[:, None], delta[None, :], deltas) + outputs = projections + tl.dot( + coefficients.to(projection_dtype), deltas.to(projection_dtype), input_precision=projection_precision + ) + out_tokens = tl.where(valid, tokens, start) + tl.store( + Out + (out_tokens[:, None] * HV + hv) * V + vv[None, :], + outputs, + (valid & (tt < BW))[:, None] & (vv[None, :] < V), + ) + records = n + jj - L + in_prefix = fresh & (records < L) + safe_records = tl.where(in_prefix, records, 0) + tl.store( + Deltas + (slot * L + safe_records[:, None]) * V + vv[None, :], + deltas, + in_prefix[:, None] & (vv[None, :] < V), + ) + tl.store( + RawValues + (slot * L + safe_records[:, None]) * V + vv[None, :], + values, + in_prefix[:, None] & (vv[None, :] < V), + ) + if iv == 0: + if hv % (HV // KH) == 0: + tl.store( + Keys + (key_slot * L + safe_records[:, None]) * K + kk[None, :], + fresh_keys, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store( + RawKeys + (key_slot * L + safe_records[:, None]) * K + kk[None, :], + raw_keys, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store(Gates + slot * L + safe_records, fresh_gates, in_prefix) + tl.store(Betas + slot * L + safe_records, betas, in_prefix) + return + + +@triton.jit +def _kda_precomputed_window( + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SK: tl.constexpr, + SV: tl.constexpr, + SA: tl.constexpr, + SB: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + BV: tl.constexpr, + BW: tl.constexpr, + BR: tl.constexpr, + BK: tl.constexpr, + LOWER: tl.constexpr, + end, + h, + hv, + iv, + key_slot, + kk, + n, + projections, + qk, + slot, + start, + tt, + vv, +): + # KDA 的衰减在每个 key 通道上，不能把标量 exp 提到点积外面。 + # 每个 token 用冻结 checkpoint 的投影，加上 history 的通道衰减点积重建输出。 + # exp(G_t - G_j) 对已可见的 token 不大于 1。新 delta 仍按 token 解。 + # 状态压缩继续用原始 k/v 在 fp32 里折回。 + jj = tl.arange(0, BR) + fresh_tokens = start + jj - L + old = (jj < L) & (jj < n) + fresh = (jj >= L) & (jj < L + end - start) + safe_jj = tl.where(jj < L, jj, 0) + safe_fresh = tl.where(fresh, fresh_tokens, start) + cached_keys = tl.load( + Keys + (key_slot * L) * K + safe_jj[:, None] * K + kk[None, :], + old[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + raw_keys = tl.load( + Kp + safe_fresh[:, None] * SK + h * K + kk[None, :], + fresh[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + fresh_keys = raw_keys * tl.rsqrt(tl.sum(raw_keys * raw_keys, axis=1) + 1.0e-6)[:, None] + fresh_keys = tl.where(fresh[:, None], fresh_keys, 0) + all_keys = tl.where(old[:, None], cached_keys, 0) + fresh_keys + values = tl.load( + Vp + safe_fresh[:, None] * SV + hv * V + vv[None, :], + fresh[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + hist_g = tl.load( + Gates + (slot * L + safe_jj[:, None]) * K + kk[None, :], + old[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + total_hist = tl.sum(hist_g, axis=0) + hist_prefix = tl.cumsum(hist_g, axis=0) + wi = tl.arange(0, BW) + gate_ok = start + wi < end + safe_wi = tl.where(gate_ok, wi, 0) + gate_inputs = tl.load( + A + (start + safe_wi[:, None]) * SA + hv * K + kk[None, :], + gate_ok[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + gate_inputs += tl.load(Bias + hv * K + kk, kk < K, 0).to(tl.float32)[None, :] + fresh_g_rows = LOWER * tl.sigmoid(tl.exp(tl.load(Alog + hv).to(tl.float32)) * gate_inputs) + fresh_g_rows = tl.where(gate_ok[:, None], fresh_g_rows, 0) + fresh_prefix = tl.cumsum(fresh_g_rows, axis=0) + fresh_index = tl.broadcast_to(tl.where(fresh, jj - L, 0)[:, None], (BR, BK)) + G_pos = tl.where(old[:, None], hist_prefix, 0.0) + G_pos = tl.where(fresh[:, None], total_hist[None, :] + tl.gather(fresh_prefix, fresh_index, 0), G_pos) + step_g = tl.where(fresh[:, None], tl.gather(fresh_g_rows, fresh_index, 0), 0.0) + G_query = total_hist[None, :] + tl.gather(fresh_prefix, tl.broadcast_to((tt % BW)[:, None], qk.shape), 0) + deltas = tl.load( + Deltas + (slot * L) * V + safe_jj[:, None] * V + vv[None, :], + old[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + betas = tl.sigmoid(tl.load(B + safe_fresh * SB + hv, fresh, 0).to(tl.float32)) + for t in range(start, end): + local = t - start + q = tl.sum(tl.where((tt == local)[:, None], qk, 0), 0) + k = tl.sum(tl.where((tt == BW + local)[:, None], qk, 0), 0) + g_now = tl.sum(tl.where((tt == local)[:, None], G_query, 0), 0) + causal = old | (fresh & (jj <= L + local)) + decay = tl.exp(tl.where(causal[:, None], g_now[None, :] - G_pos, 0.0)) + scaled = tl.where(causal[:, None], all_keys * decay, 0) + hk = _history_coeff(scaled, k, BK) + proj_k = tl.sum(tl.where((tt == BW + local)[:, None], projections, 0), 0) + sk = proj_k + _apply_deltas(hk, deltas, BR, BV) + value = tl.sum(tl.where((jj == L + local)[:, None], values, 0), 0) + beta = tl.sum(tl.where(jj == L + local, betas, 0)) + delta = beta * (value - sk) + deltas = tl.where((jj == L + local)[:, None], delta[None, :], deltas) + hq = _history_coeff(scaled, q, BK) + proj_q = tl.sum(tl.where((tt == local)[:, None], projections, 0), 0) + sq = proj_q + _apply_deltas(hq, deltas, BR, BV) + tl.store(Out + (t * HV + hv) * V + vv, sq, vv < V) + records = n + jj - L + in_prefix = fresh & (records < L) + safe_records = tl.where(in_prefix, records, 0) + tl.store( + Deltas + (slot * L + safe_records[:, None]) * V + vv[None, :], + deltas, + in_prefix[:, None] & (vv[None, :] < V), + ) + tl.store( + RawValues + (slot * L + safe_records[:, None]) * V + vv[None, :], + values, + in_prefix[:, None] & (vv[None, :] < V), + ) + if iv == 0: + if hv % (HV // KH) == 0: + tl.store( + Keys + (key_slot * L + safe_records[:, None]) * K + kk[None, :], + fresh_keys, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store( + RawKeys + (key_slot * L + safe_records[:, None]) * K + kk[None, :], + raw_keys, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store( + Gates + (slot * L + safe_records[:, None]) * K + kk[None, :], + step_g, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store(Betas + slot * L + safe_records, betas, in_prefix) + return + + +@triton.jit +def _replay_each_token( + Q, + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SQ: tl.constexpr, + SK: tl.constexpr, + SV: tl.constexpr, + SA: tl.constexpr, + SB: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + PRECOMPUTE_PROJECTIONS: tl.constexpr, + BW: tl.constexpr, + KDA: tl.constexpr, + LOWER: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + end, + hv, + iv, + key_slot, + kk, + ll, + n, + projections, + slot, + start, + state, + tt, + vv, +): + # KDA、inline 以及不满足 precompute 条件的路径使用逐 token 递推。 + history = ll < n + keys = tl.load( + Keys + (key_slot * L) * K + ll[:, None] * K + kk[None, :], + history[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + deltas = tl.load( + Deltas + (slot * L) * V + ll[:, None] * V + vv[None, :], + history[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + if KDA: + gates = tl.load( + Gates + (slot * L + ll[:, None]) * K + kk[None, :], + history[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + else: + gates = tl.load(Gates + slot * L + ll, history, 0).to(tl.float32) + gate_prefix = tl.cumsum(gates, axis=0) + total_g = tl.sum(gates, axis=0) + if KDA: + keys *= tl.where(history[:, None], tl.exp(total_g[None, :] - gate_prefix), 0.0) + h = hv // (HV // H) + log_a = tl.load(Alog + hv).to(tl.float32) + bias = tl.load(Bias + hv * K + kk, kk < K, 0).to(tl.float32) if KDA else tl.load(Bias + hv).to(tl.float32) + + for t in range(start, end): + q = tl.load(Q + t * SQ + h * K + kk, kk < K, 0).to(tl.float32) + raw_k = tl.load(Kp + t * SK + h * K + kk, kk < K, 0).to(tl.float32) + v = tl.load(Vp + t * SV + hv * V + vv, vv < V, 0).to(tl.float32) + q = q / tl.sqrt(tl.sum(q * q) + 1.0e-6) * (K ** -0.5) + k = raw_k / tl.sqrt(tl.sum(raw_k * raw_k) + 1.0e-6) + if KDA: + x = tl.load(A + t * SA + hv * K + kk, kk < K, 0).to(tl.float32) + bias + g = LOWER * tl.sigmoid(tl.exp(log_a) * x) + else: + x = tl.load(A + t * SA + hv).to(tl.float32) + bias + g = -tl.exp(log_a) * tl.where(x <= 20.0, tl.log(1.0 + tl.exp(x)), x) + beta = tl.sigmoid(tl.load(B + t * SB + hv).to(tl.float32)) + total_g += g + if KDA: + # Advance the decayed keys once per token, rather than + # evaluating L independent vector exponentials each time. + keys *= tl.exp(g)[None, :] + hk = _history_coeff(keys, k, BK) + hq = _history_coeff(keys, q, BK) + if PRECOMPUTE_PROJECTIONS: + checkpoint_sk = tl.sum(tl.where((tt == BW + t - start)[:, None], projections, 0), 0) + checkpoint_sq = tl.sum(tl.where((tt == t - start)[:, None], projections, 0), 0) + else: + checkpoint_sk = tl.sum(state * (k * tl.exp(total_g))[:, None], axis=0) + checkpoint_sq = tl.sum(state * (q * tl.exp(total_g))[:, None], axis=0) + sk = checkpoint_sk + _apply_deltas(hk, deltas, L, BV) + sq = checkpoint_sq + _apply_deltas(hq, deltas, L, BV) + else: + weights = tl.where(history, tl.exp(total_g - gate_prefix), 0.0) + hk = _history_coeff(keys, k, BK) * weights + hq = _history_coeff(keys, q, BK) * weights + checkpoint_sk = tl.sum(state * k[:, None], axis=0) + checkpoint_sq = tl.sum(state * q[:, None], axis=0) + sk = checkpoint_sk * tl.exp(total_g) + _apply_deltas(hk, deltas, L, BV) + sq = checkpoint_sq * tl.exp(total_g) + _apply_deltas(hq, deltas, L, BV) + d = beta * (v - sk) + tl.store(Out + (t * HV + hv) * V + vv, sq + d * tl.sum(k * q), vv < V) + record = n + in_prefix = record < L + # 被 mask 掉的 lane 仍会计算地址。cursor 已经顶到容量时必须改写 0 号槽，避免非法访问。 + safe_record = tl.where(in_prefix, record, 0) + tl.store(Deltas + (slot * L + safe_record) * V + vv, d, in_prefix & (vv < V)) + tl.store(RawValues + (slot * L + safe_record) * V + vv, v, in_prefix & (vv < V)) + if iv == 0: + # Only the group owner writes shared K. Every CTA computes this + # round's K locally; readers only consume earlier accepted records. + if hv % (HV // KH) == 0: + tl.store(Keys + (key_slot * L + safe_record) * K + kk, k, in_prefix & (kk < K)) + tl.store(RawKeys + (key_slot * L + safe_record) * K + kk, raw_k, in_prefix & (kk < K)) + if KDA: + tl.store(Gates + (slot * L + safe_record) * K + kk, g, in_prefix & (kk < K)) + else: + tl.store(Gates + slot * L + safe_record, g, in_prefix) + tl.store(Betas + slot * L + safe_record, beta, in_prefix) + keys = tl.where((ll == n)[:, None], k[None, :], keys) + deltas = tl.where((ll == n)[:, None], d[None, :], deltas) + if not KDA: + gate_prefix = tl.where(ll >= n, total_g, gate_prefix) + history |= ll == n + n += 1 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py new file mode 100644 index 0000000000..b3a6b2aec0 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +"""One gated-delta step shared by prefix folding.""" + +import triton +import triton.language as tl + + +@triton.jit +def _delta_rule_step(state, raw_k, raw_v, g, beta, KDA: tl.constexpr): + # Folding applies this same update once per accepted history token. + raw_k /= tl.sqrt(tl.sum(raw_k * raw_k) + 1.0e-6) + decay = tl.exp(g) + if KDA: + state *= decay[:, None] + d = beta * (raw_v - tl.sum(state * raw_k[:, None], 0)) + state += raw_k[:, None] * d[None, :] + else: + sk = tl.sum(state * raw_k[:, None], 0) + d = beta * (raw_v - decay * sk) + state = state * decay + raw_k[:, None] * d[None, :] + return state diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py new file mode 100644 index 0000000000..7ac6a25c8b --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fold an accepted prefix into the canonical SSM.""" + +import triton +import triton.language as tl + +from .replayssm_delta import _delta_rule_step + + +@triton.jit(do_not_specialize=["N"]) +def _reset_folded_cursors( + Cursors, Reqs, N, HOLD: tl.constexpr, L: tl.constexpr, NEXT: tl.constexpr, BLOCK: tl.constexpr +): + row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + req = tl.load(Reqs + row, row < N, HOLD) + valid = (row < N) & (req != HOLD) + n = tl.load(Cursors + req, valid, 0) + tl.store(Cursors + req, 0, valid & (n + NEXT > L)) + + +@triton.jit +def _fold_prefix( + State, + Output, + RawKeys, + RawValues, + Gates, + Betas, + Cursors, + Reqs, + STATE_STRIDES: tl.constexpr, + SLOTS: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + HOLD: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + SNAPSHOT: tl.constexpr, + KDA: tl.constexpr, + FLUSH_EXPORT: tl.constexpr = False, + OUTPUT_STRIDES: tl.constexpr = None, + WHEN_FULL: tl.constexpr = False, + NEXT: tl.constexpr = 1, +): + # 将 cursor 指向的 accepted history 应用到 canonical state。SNAPSHOT + # 只写独立输出。WHEN_FULL 只压缩「长度 + 下一次最大宽度」放不下的请求， + # cursor 在后续 kernel 中清零；这里的所有 layer/head tile 都需要读到原始长度。 + iv, row, lh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + req = tl.load(Reqs + row).to(tl.int64) + if req == HOLD: + return + n = tl.load(Cursors + req) + if WHEN_FULL: + if n == 0 or n + NEXT <= L: + return + elif n == 0 and not SNAPSHOT and not FLUSH_EXPORT: + return + linear_layer_idx, head = lh // HV, lh % HV + slot = (linear_layer_idx * SLOTS + req) * HV + head + key_slot = (linear_layer_idx * SLOTS + req) * KH + head // (HV // KH) + kk = tl.arange(0, BK) + vv = iv * BV + tl.arange(0, BV) + state_offset = ( + linear_layer_idx.to(tl.int64) * STATE_STRIDES[0] + + req * STATE_STRIDES[1] + + head.to(tl.int64) * STATE_STRIDES[2] + + kk[:, None].to(tl.int64) * STATE_STRIDES[3] + + vv[None, :].to(tl.int64) * STATE_STRIDES[4] + ) + sp = State + state_offset + state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) + # cursor 可能暂时超过容量。循环只走容量内的槽，越界步不读 history，也不更新 state。 + for j in range(L): + active = j < n + raw_k = tl.load(RawKeys + (key_slot * L + j) * K + kk, active & (kk < K), 0).to(tl.float32) + raw_v = tl.load(RawValues + (slot * L + j) * V + vv, active & (vv < V), 0).to(tl.float32) + if KDA: + g = tl.load(Gates + (slot * L + j) * K + kk, active & (kk < K), 0).to(tl.float32) + else: + g = tl.load(Gates + slot * L + j, active, 0).to(tl.float32) + beta = tl.load(Betas + slot * L + j, active, 0).to(tl.float32) + updated = _delta_rule_step(state, raw_k, raw_v, g, beta, KDA) + state = tl.where(active, updated, state) + mask = (kk[:, None] < K) & (vv[None, :] < V) + if OUTPUT_STRIDES is not None: + target = tl.load(Output + row).to(tl.pointer_type(State.dtype.element_ty)) + out = ( + target + + linear_layer_idx.to(tl.int64) * OUTPUT_STRIDES[0] + + head * OUTPUT_STRIDES[1] + + kk[:, None] * OUTPUT_STRIDES[2] + + vv[None, :] * OUTPUT_STRIDES[3] + ) + else: + if SNAPSHOT: + out = Output + (linear_layer_idx * HV + head) * K * V + kk[:, None] * V + vv[None, :] + else: + out = Output + state_offset + tl.store(out, state, mask) + if FLUSH_EXPORT: + tl.store(sp, state, mask) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py new file mode 100644 index 0000000000..5244dada83 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +def commit_ssm_prefix( + indices: torch.Tensor, + reqs: torch.Tensor, + lengths: torch.Tensor, + replay_cursors: torch.Tensor | None = None, +) -> None: + """把每个请求的 accepted length 写入 MTP 索引并推进 Replay cursor。""" + if not reqs.is_cuda: + indices[reqs] = (lengths - 1).to(indices.dtype) + if replay_cursors is not None: + valid = reqs != replay_cursors.numel() - 1 + replay_cursors[reqs[valid]] += lengths[valid].to(replay_cursors.dtype) + return + if reqs.numel(): + _commit_prefix[(triton.cdiv(reqs.numel(), 256),)]( + indices, + reqs, + lengths, + replay_cursors, + reqs.numel(), + indices.numel() - 1, + reqs.stride(0), + lengths.stride(0), + 256, + num_warps=4, + ) + + +@triton.jit(do_not_specialize=["N"]) +def _commit_prefix( + Indices, + Reqs, + Lengths, + Cursors, + N, + HOLD: tl.constexpr, + REQ_STRIDE: tl.constexpr, + LENGTH_STRIDE: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + req = tl.load(Reqs + row * REQ_STRIDE, row < N, 0) + length = tl.load(Lengths + row * LENGTH_STRIDE, row < N, 0) + tl.store(Indices + req, length - 1, row < N) + if Cursors is not None: + valid = (row < N) & (req != HOLD) + cursor = tl.load(Cursors + req, valid, 0) + tl.store(Cursors + req, cursor + length, valid) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py new file mode 100644 index 0000000000..7ca84503c1 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -0,0 +1,180 @@ +"""使用共享 autotuner 调优单层 SSM update cycle。""" + +from __future__ import annotations + +import ast + +import torch +import triton + +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, autotune +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + +_DEFAULT_REPLAY_CONFIG = {"BV": 32, "num_warps": 1, "num_stages": 3} + + +def resolve_run_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens) -> dict[str, int]: + """返回当前 batch 的固定或已调优 SSM kernel 配置。""" + if ssm_backend._config_is_fixed and ssm_backend.run_config is not None: + return ssm_backend.run_config + default = _replay_config(ssm_backend.run_config) + key = (tuple(static_key(ssm_backend, mode, q, cu_seqlens).items()), run_key(q, cu_seqlens)) + if key in ssm_backend.batch_configs: + return ssm_backend.batch_configs[key] + if not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + return default + config = _replay_config(select_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens)) + # 每个 captured graph 保留自己的布局，即使其他 bucket 后续完成 warmup。 + ssm_backend.batch_configs[key] = config + logger.info(f"SSM {mode} update-cycle config: {config}, warmup batch: {key[1]}") + return config + + +def get_configs() -> list[dict[str, int]]: + """返回 ReplaySSM update cycle 的候选 Triton 配置。""" + return [ + {"BV": bv, "num_warps": warps, "num_stages": stages} + for bv, warps in [(8, 1), (16, 1), (16, 2), (32, 1), (32, 2), (32, 4), (64, 2), (64, 4), (128, 4)] + for stages in (1, 2, 3) + ] + + +def static_key(ssm_backend, mode, q, cu_seqlens) -> dict[str, object]: + """构造区分状态布局、数据类型和投影模式的静态调优键。""" + _, _, hv, kd, vd = ssm_backend.ssm_state.shape + # 保持文件名较短，同时区分 packed projection、TP 布局、状态精度、递推规则和 history 容量。 + key = { + "mode": mode, + "v": 7, + "shape": str((q.shape[-2], hv, ssm_backend.num_key_heads, kd, vd, ssm_backend.verify_width)), + "dtype": str((q.dtype, ssm_backend.ssm_state.dtype)), + "history": ssm_backend.capacity, + "gate": int(ssm_backend.kda), + "varlen": cu_seqlens is not None, + } + if ssm_backend.kda: + key["lower"] = ssm_backend.gate_lower_bound + key["history_dtype"] = str(ssm_backend.keys.dtype) + key["delta_dtype"] = str(ssm_backend.deltas.dtype) + key["projection"] = ( + ("kda_mma" if ssm_backend.kda else "gdn_mma") if ssm_backend.compute_mode == "precompute" else "inline" + ) + return key + + +def run_key(q, cu_seqlens) -> tuple[int, int]: + """返回调优使用的 token 数和请求数键。""" + tokens = q.shape[1 if cu_seqlens is not None else 0] + return tokens, cu_seqlens.numel() - 1 if cu_seqlens is not None else tokens + + +def rebuild_inputs( + ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None +) -> tuple[tuple[object, ...], dict[str, object]]: + """为一层 scratch 和同形状的紧凑激活分配调优输入。不复制服务状态池。""" + tokens, sequences = run_key(q, cu_seqlens) + batch = min(sequences, triton.cdiv(tokens, ssm_backend.verify_width)) if cu_seqlens is not None else tokens + _, _, hv, kd, vd = ssm_backend.ssm_state.shape + from lightllm.common.attention_state.bounded_att.conv_ssm_types import ConvSSMStateConfig + + # 计时只需要一层、本 batch 的槽位。正式执行仍用调用方原来的 state。 + scratch = type(ssm_backend)( + ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=ssm_backend.num_key_heads, + global_v_heads=hv, + key_dim=kd, + value_dim=vd, + conv_width=1, + conv_dtype=q.dtype, + ssm_dtype=ssm_backend.ssm_state.dtype, + ), + batch + 1, + mtp_step=ssm_backend.verify_width - 1, + device=q.device, + ssm_replay_config={"cache_len": ssm_backend.capacity, "compute_mode": ssm_backend.compute_mode}, + gate_lower_bound=ssm_backend.gate_lower_bound, + ) + scratch.ssm_state.normal_(0, 0.02) + # benchmark 回调不能再次触发配置选择。 + scratch._config_is_fixed = True + inputs = [ + torch.empty(x.shape, device=x.device, dtype=x.dtype).normal_(0, 0.2) for x in (q, k, v, a, b, a_log, bias) + ] + reqs = torch.cat( + ( + torch.arange(batch, dtype=torch.int32, device=q.device), + torch.full((sequences - batch,), batch, dtype=torch.int32, device=q.device), + ) + ) + cu = (torch.arange(sequences + 1, dtype=torch.int32, device=q.device) * ssm_backend.verify_width).clamp_max(tokens) + lengths = cu[1 : batch + 1] - cu[:batch] + accepted = torch.cat((lengths - 1, lengths.new_zeros(1))) + return (), dict( + ssm_backend=scratch, + mode=mode, + q=inputs[0], + k=inputs[1], + v=inputs[2], + a=inputs[3], + b=inputs[4], + a_log=inputs[5], + bias=inputs[6], + cu_seqlens=cu if cu_seqlens is not None else None, + workload=(reqs, accepted), + ) + + +def replay_layer(history, *args, **kwargs): + """按 history 的具体类型调用对应的 replay decode。""" + + 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 + + if isinstance(history, ReplayKDAState): + from lightllm.common.basemodel.attention.linear.replay_kda_decode import replay + elif isinstance(history, ReplayGDNState): + from lightllm.common.basemodel.attention.linear.replay_gdn_decode import replay + else: + raise TypeError(f"replay_layer expects ReplayGDNState or ReplayKDAState, got {type(history).__name__}") + return replay(history, *args, **kwargs) + + +@autotune( + kernel_name="ssm_update_cycle:v2", + configs_gen_func=get_configs, + static_key_func=static_key, + run_key_func=run_key, + run_key_distance_func=lambda a, b: sum(abs(x - y) for x, y in zip(ast.literal_eval(a), ast.literal_eval(b))), + kernel_type=AutotuneKernelType.DECODE_ATTENTION, + rebuild_input_func=rebuild_inputs, + warmup_all_exist_config=False, +) +def select_config( + ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None +) -> dict[str, int] | None: + """搜索 update cycle 布局；计时只使用重建的临时输入。""" + if workload is None: + return run_config + ssm_backend.run_config = run_config + reqs, accepted = workload + ssm_backend.cursors.zero_() + # 把前缀写满但不越出 capacity。replay 自己读 cursor，正式 decode 也不经过 prepare_decode。 + tokens = q.shape[1 if cu_seqlens is not None else 0] + rounds = max(1, ssm_backend.capacity // min(ssm_backend.verify_width, tokens)) + for _ in range(rounds): + replay_layer(ssm_backend, 0, q, k, v, a, b, a_log, bias, reqs, cu_seqlens) + ssm_backend.accept_updates(reqs, accepted) + return run_config + + +def _replay_config(config: dict | None) -> dict[str, int]: + if not config: + return dict(_DEFAULT_REPLAY_CONFIG) + return { + "BV": int(config["BV"]), + "num_warps": int(config["num_warps"]), + "num_stages": int(config.get("num_stages", 3)), + } diff --git a/lightllm/common/basemodel/triton_kernel/mtp_utils.py b/lightllm/common/basemodel/triton_kernel/mtp_utils.py index 7525a754df..85ac266691 100644 --- a/lightllm/common/basemodel/triton_kernel/mtp_utils.py +++ b/lightllm/common/basemodel/triton_kernel/mtp_utils.py @@ -309,12 +309,11 @@ def mtp_scatter_next_token_ids( ) -@triton.jit +@triton.jit(do_not_specialize=["batch_size"]) def _fwd_kernel_gen_b_req_mtp_start_loc( b_mtp_index, b_req_mtp_start_loc, - num_reqs: tl.constexpr, - batch_size: tl.constexpr, + batch_size, BLOCK_SIZE: tl.constexpr, ): offset = tl.arange(0, BLOCK_SIZE) @@ -333,7 +332,6 @@ def gen_b_req_mtp_start_loc(b_mtp_index: torch.Tensor, num_reqs: int): _fwd_kernel_gen_b_req_mtp_start_loc[grid]( b_mtp_index=b_mtp_index, b_req_mtp_start_loc=b_req_mtp_start_loc, - num_reqs=num_reqs, batch_size=batch_size, BLOCK_SIZE=BLOCK_SIZE, num_warps=8, diff --git a/lightllm/common/basemodel/triton_kernel/scalar_linear.py b/lightllm/common/basemodel/triton_kernel/scalar_linear.py new file mode 100644 index 0000000000..e78678c40a --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/scalar_linear.py @@ -0,0 +1,42 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _scalar_linear(X, W, Y, K: tl.constexpr, XS0, XS1, WS, YS, BLOCK: tl.constexpr): + row = tl.program_id(0) + column = tl.arange(0, BLOCK) + value = tl.load(X + row * XS0 + column * XS1, column < K, other=0).to(tl.float64) + weight = tl.load(W + column * WS, column < K, other=0).to(tl.float64) + tl.store(Y + row * YS, tl.sum(value * weight, 0)) + + +def scalar_linear(input: torch.Tensor, weight: torch.Tensor, out=None, alloc_func=torch.empty) -> torch.Tensor: + if input.ndim != 2 or weight.shape != (1, input.shape[1]): + raise ValueError("Scalar projection requires (tokens, hidden) input and (1, hidden) weight") + if not input.is_cuda or weight.device != input.device or weight.dtype != input.dtype: + raise ValueError("Scalar projection requires CUDA input and matching weight dtype/device") + if not input.is_floating_point(): + raise ValueError("Scalar projection requires floating-point activations") + shape = (input.shape[0], 1) + if out is None: + out = alloc_func(shape, dtype=input.dtype, device=input.device) + elif out.shape != shape or out.dtype != input.dtype or out.device != input.device: + raise ValueError("Scalar projection output must match the unpadded shape, dtype and device") + if input.shape[0]: + block = triton.next_power_of_2(max(1, input.shape[1])) + _scalar_linear[(input.shape[0],)]( + input, + weight, + out, + input.shape[1], + input.stride(0), + input.stride(1), + weight.stride(1), + out.stride(0), + block, + num_warps=4 if block <= 8192 else 8, + enable_fp_fusion=False, + ) + return out diff --git a/lightllm/common/gated_residual/__init__.py b/lightllm/common/gated_residual/__init__.py new file mode 100644 index 0000000000..7f0da1f957 --- /dev/null +++ b/lightllm/common/gated_residual/__init__.py @@ -0,0 +1,85 @@ +from dataclasses import dataclass +from typing import Mapping, Optional + +import torch + +from .kernels import grouped_norm, inject_streams, injection_gates +from .projection import project_and_mix + + +@dataclass(frozen=True) +class GatedResidual: + norm_weight: torch.Tensor + down_weight: torch.Tensor + up_weight: torch.Tensor + inject_weight: Optional[torch.Tensor] + branches: int + eps: float = 1e-6 + + def __post_init__(self): + width = self.norm_weight.numel() + if self.branches < 2 or width == 0 or width % self.branches: + raise ValueError("Residual width must be a positive multiple of branches >= 2") + if self.norm_weight.shape != (width,) or self.down_weight.ndim != 2: + raise ValueError("Expected a norm vector and a matrix for the down projection") + rank = self.down_weight.shape[0] + if rank == 0 or self.down_weight.shape[1] != width or self.up_weight.shape != (width, rank): + raise ValueError("Residual projection dimensions do not match") + if self.inject_weight is not None and self.inject_weight.shape != (self.branches, width): + raise ValueError("Injection projection must have shape (branches, residual width)") + weights = [self.norm_weight, self.down_weight, self.up_weight] + if self.inject_weight is not None: + weights.append(self.inject_weight) + if self.eps <= 0: + raise ValueError("Normalization epsilon must be positive") + for weight in weights: + if not weight.is_cuda or not weight.is_contiguous(): + raise ValueError("Gated residual weights must be contiguous CUDA tensors") + if weight.device != self.norm_weight.device or weight.dtype != self.norm_weight.dtype: + raise ValueError("Gated residual weights must share a device and dtype") + if self.norm_weight.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError("Gated residuals require floating point weights") + + @classmethod + def from_checkpoint( + cls, + weights: Mapping[str, torch.Tensor], + prefix: str, + branches: int, + eps: float = 1e-6, + use_combine: bool = True, + ) -> "GatedResidual": + return cls( + weights[f"{prefix}.hc_norm.weight"], + weights[f"{prefix}.input_mix_weight_down.weight"], + weights[f"{prefix}.input_mix_weight_up.weight"], + weights[f"{prefix}.block_inject_weight.weight"] if use_combine else None, + branches, + eps, + ) + + def _check_input(self, x: torch.Tensor, width: int): + if x.ndim < 2 or x.shape[-1] != width or not x.is_contiguous(): + raise ValueError(f"Expected contiguous input with last dimension {width}") + if x.device != self.norm_weight.device or x.dtype != self.norm_weight.dtype: + raise ValueError("Input must share the weights' device and dtype") + + @torch.no_grad() + def read(self, residual: torch.Tensor): + self._check_input(residual, self.norm_weight.numel()) + normalized = grouped_norm(residual, self.norm_weight, self.branches, self.eps) + mixed = project_and_mix(normalized, self.down_weight, self.up_weight, self.branches) + if self.inject_weight is None: + return mixed + gates = injection_gates(normalized, self.inject_weight, self.branches) + return mixed, gates + + @torch.no_grad() + def write(self, residual: torch.Tensor, update: torch.Tensor, gates: torch.Tensor) -> torch.Tensor: + width = self.norm_weight.numel() + self._check_input(residual, width) + self._check_input(update, width // self.branches) + self._check_input(gates, self.branches) + if update.shape[:-1] != residual.shape[:-1] or gates.shape[:-1] != residual.shape[:-1]: + raise ValueError("Residual, update, and gates must have matching token dimensions") + return inject_streams(residual, update, gates) diff --git a/lightllm/common/gated_residual/kernels.py b/lightllm/common/gated_residual/kernels.py new file mode 100644 index 0000000000..96f7271db0 --- /dev/null +++ b/lightllm/common/gated_residual/kernels.py @@ -0,0 +1,103 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _norm(X, W, Y, D: tl.constexpr, C: tl.constexpr, EPS: tl.constexpr, B: tl.constexpr): + row = tl.program_id(0) + branch = tl.program_id(1) + col = tl.arange(0, B) + offset = row * C * D + branch * D + col + x = tl.load(X + offset, col < D, other=0).to(tl.float32) + weight = tl.load(W + branch * D + col, col < D, other=0).to(tl.float32) + inv = tl.rsqrt(tl.sum(x * x, 0) / D + EPS) + tl.store(Y + offset, (x * inv) * (1.0 + weight), col < D) + + +@triton.jit +def _mix(X, G, Y, D: tl.constexpr, C: tl.constexpr, BC: tl.constexpr, B: tl.constexpr): + row = tl.program_id(0) + col = tl.program_id(1) * B + tl.arange(0, B) + branch = tl.arange(0, BC) + offset = row * C * D + branch[:, None] * D + col[None, :] + mask = (branch[:, None] < C) & (col[None, :] < D) + x = tl.load(X + offset, mask, other=0).to(tl.float32) + logits = tl.load(G + offset, mask, other=0).to(tl.float32) + + gate = tl.sigmoid(logits).to(G.dtype.element_ty).to(tl.float32) + product = (x * gate).to(X.dtype.element_ty).to(tl.float32) + product = tl.where(branch[:, None] < C, product, 0.0) + result = tl.sum(product, 0) / C + tl.store(Y + row * D + col, result, col < D) + + +@triton.jit +def _gates(X, W, Out, D: tl.constexpr, C: tl.constexpr, B: tl.constexpr): + row = tl.program_id(0) + branch = tl.program_id(1) + col = tl.arange(0, B) + x = tl.load(X + row * D + col, col < D, other=0).to(tl.float32) + weight = tl.load(W + branch * D + col, col < D, other=0).to(tl.float32) + + logits = tl.sum(x * weight, 0).to(X.dtype.element_ty).to(tl.float32) + scaled = (logits / C).to(X.dtype.element_ty).to(tl.float32) + gate = tl.sigmoid(scaled).to(X.dtype.element_ty).to(tl.float32) + tl.store(Out + row * C + branch, 2 * gate) + + +@triton.jit +def _inject(R, X, G, Y, D: tl.constexpr, C: tl.constexpr, B: tl.constexpr): + row = tl.program_id(0) + branch = tl.program_id(1) + col = tl.program_id(2) * B + tl.arange(0, B) + offset = row * C * D + branch * D + col + x = tl.load(X + row * D + col, col < D, other=0).to(tl.float32) + gate = tl.load(G + row * C + branch).to(tl.float32) + residual = tl.load(R + offset, col < D, other=0).to(tl.float32) + + injection = (x * gate).to(R.dtype.element_ty).to(tl.float32) + tl.store(Y + offset, residual + injection, col < D) + + +def grouped_norm(x: torch.Tensor, weight: torch.Tensor, branches: int, eps: float) -> torch.Tensor: + output = torch.empty_like(x) + rows = x.numel() // x.shape[-1] + dim = x.shape[-1] // branches + if rows: + _norm[(rows, branches)]( + x, weight, output, dim, branches, eps, triton.next_power_of_2(dim), enable_fp_fusion=False + ) + return output + + +def mix_streams(x: torch.Tensor, logits: torch.Tensor, branches: int) -> torch.Tensor: + dim = x.shape[-1] // branches + output = torch.empty((*x.shape[:-1], dim), dtype=x.dtype, device=x.device) + rows = x.numel() // x.shape[-1] + if rows: + _mix[(rows, triton.cdiv(dim, 256))](x, logits, output, dim, branches, triton.next_power_of_2(branches), 256) + return output + + +def inject_streams(residual: torch.Tensor, update: torch.Tensor, gates: torch.Tensor) -> torch.Tensor: + output = torch.empty_like(residual) + branches = gates.shape[-1] + dim = update.shape[-1] + rows = update.numel() // dim + if rows: + _inject[(rows, branches, triton.cdiv(dim, 256))]( + residual, update, gates, output, dim, branches, 256, enable_fp_fusion=False + ) + return output + + +def injection_gates(x: torch.Tensor, weight: torch.Tensor, branches: int) -> torch.Tensor: + output = torch.empty((*x.shape[:-1], branches), dtype=x.dtype, device=x.device) + rows = x.numel() // x.shape[-1] + if rows: + block = triton.next_power_of_2(x.shape[-1]) + _gates[(rows, branches)]( + x, weight, output, x.shape[-1], branches, block, num_warps=8 if block >= 8192 else 4, enable_fp_fusion=False + ) + return output diff --git a/lightllm/common/gated_residual/projection.py b/lightllm/common/gated_residual/projection.py new file mode 100644 index 0000000000..8d498650b1 --- /dev/null +++ b/lightllm/common/gated_residual/projection.py @@ -0,0 +1,156 @@ +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from .kernels import mix_streams + + +@triton.jit +def _activate(value, DTYPE: tl.constexpr, C: tl.constexpr): + value = value.to(DTYPE).to(tl.float32) + value = (value / C).to(DTYPE).to(tl.float32) + return value * tl.sigmoid(value) + + +@triton.jit +def _down( + X, + W, + P, + Y, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, + C: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + S: tl.constexpr, +): + r = tl.program_id(0) * BM + tl.arange(0, BM) + c = tl.program_id(1) * BN + tl.arange(0, BN) + split = tl.program_id(2) + part = tl.cdiv(K, BK * S) * BK + k = split * part + tl.arange(0, BK) + acc = tl.full((BM, BN), 0, tl.float32) + for _ in range(tl.cdiv(K, BK * S)): + x = tl.load(X + r[:, None] * K + k[None, :], (r[:, None] < M) & (k[None, :] < K), other=0) + w = tl.load(W + c[None, :] * K + k[:, None], (c[None, :] < N) & (k[:, None] < K), other=0) + acc = tl.dot(x, w, acc, input_precision="ieee") + k += BK + if S == 1: + tl.store( + Y + r[:, None] * N + c[None, :], _activate(acc, Y.dtype.element_ty, C), (r[:, None] < M) & (c[None, :] < N) + ) + else: + tl.store(P + (split * M + r[:, None]) * N + c[None, :], acc, (r[:, None] < M) & (c[None, :] < N)) + + +@triton.jit +def _merge(P, Y, M: tl.constexpr, N: tl.constexpr, C: tl.constexpr, S: tl.constexpr, B: tl.constexpr): + row = tl.program_id(0) + col = tl.program_id(1) * B + tl.arange(0, B) + split = tl.arange(0, S) + value = tl.load(P + (split[:, None] * M + row) * N + col[None, :], col[None, :] < N, other=0) + tl.store(Y + row * N + col, _activate(tl.sum(value, 0), Y.dtype.element_ty, C), col < N) + + +@triton.jit +def _up_mix( + X, + W, + Normalized, + Y, + M: tl.constexpr, + D: tl.constexpr, + K: tl.constexpr, + C: tl.constexpr, + BC: tl.constexpr, + BM: tl.constexpr, + BD: tl.constexpr, + BK: tl.constexpr, +): + row = tl.program_id(0) * BM + tl.arange(0, BM) + n = tl.arange(0, BC * BD) + branch = n // BD + col = tl.program_id(1) * BD + n % BD + weight_row = branch * D + col + k = tl.arange(0, BK) + acc = tl.full((BM, BC * BD), 0, tl.float32) + for _ in range(tl.cdiv(K, BK)): + x = tl.load(X + row[:, None] * K + k[None, :], (row[:, None] < M) & (k[None, :] < K), other=0) + w = tl.load( + W + weight_row[None, :] * K + k[:, None], + (branch[None, :] < C) & (col[None, :] < D) & (k[:, None] < K), + other=0, + ) + acc = tl.dot(x, w, acc, input_precision="ieee") + k += BK + logits = acc.to(X.dtype.element_ty).to(tl.float32) + gate = tl.sigmoid(logits).to(X.dtype.element_ty).to(tl.float32) + normalized = tl.load( + Normalized + row[:, None] * (C * D) + weight_row[None, :], + (row[:, None] < M) & (branch[None, :] < C) & (col[None, :] < D), + other=0, + ).to(tl.float32) + products = (normalized * gate).to(X.dtype.element_ty).to(tl.float32) + mixed = tl.sum(products.reshape(BM, BC, BD), 1) / C + out_col = tl.program_id(1) * BD + tl.arange(0, BD) + tl.store(Y + row[:, None] * D + out_col[None, :], mixed, (row[:, None] < M) & (out_col[None, :] < D)) + + +def _lowrank(x, w, branches): + m, k = x.shape + n = w.shape[0] + out = torch.empty((m, n), device=x.device, dtype=x.dtype) + if not m: + return out + + splits = 16 if k >= 8192 and n <= 512 else 1 + p = torch.empty((splits, m, n), device=x.device, dtype=torch.float32) if splits > 1 else out + bm = 16 if m <= 16 else 32 if m <= 64 else 64 if m < 128 else 128 + _down[(triton.cdiv(m, bm), triton.cdiv(n, 64), splits)]( + x, w, p, out, m, n, k, branches, bm, 64, 64, splits, num_warps=8 if bm == 128 else 4, enable_fp_fusion=False + ) + if splits > 1: + _merge[(m, triton.cdiv(n, 128))](p, out, m, n, branches, splits, 128, num_warps=4, enable_fp_fusion=False) + return out + + +def _mix(normalized, low, w, branches): + m, k = low.shape + dim = normalized.shape[1] // branches + out = torch.empty((m, dim), device=low.device, dtype=low.dtype) + if not m: + return out + bm = 16 if m <= 16 else 32 if m <= 64 else 64 if m < 128 else 128 + _up_mix[(triton.cdiv(m, bm), triton.cdiv(dim, 32))]( + low, + w, + normalized, + out, + m, + dim, + k, + branches, + triton.next_power_of_2(branches), + bm, + 32, + 64, + num_warps=8 if bm == 128 else 4, + enable_fp_fusion=False, + ) + return out + + +def project_and_mix( + normalized: torch.Tensor, down_weight: torch.Tensor, up_weight: torch.Tensor, branches: int +) -> torch.Tensor: + if normalized.dtype == torch.float32: + low = F.silu(F.linear(normalized, down_weight) / branches) + return mix_streams(normalized, F.linear(low, up_weight), branches) + rows = normalized.reshape(-1, normalized.shape[-1]) + low = _lowrank(rows, down_weight, branches) + mixed = _mix(rows, low, up_weight, branches) + return mixed.reshape((*normalized.shape[:-1], normalized.shape[-1] // branches)) diff --git a/lightllm/common/ngram_embedding/__init__.py b/lightllm/common/ngram_embedding/__init__.py new file mode 100644 index 0000000000..9b10a8e5a1 --- /dev/null +++ b/lightllm/common/ngram_embedding/__init__.py @@ -0,0 +1,132 @@ +from dataclasses import dataclass +from typing import Mapping, Optional + +import torch +import triton + +from .kernels import _hash, _history + + +@dataclass(frozen=True) +class NGramHasher: + multipliers: torch.Tensor + head_vocab_sizes: torch.Tensor + head_offsets: torch.Tensor + eos_token_id: int + + def __post_init__(self): + if self.multipliers.ndim != 1 or self.multipliers.numel() < 2: + raise ValueError("N-gram hashing requires at least two multipliers") + heads = self.head_vocab_sizes.numel() + if heads == 0 or heads % self.context_len or self.head_vocab_sizes.shape != (heads,): + raise ValueError("Head count must be a positive multiple of ngram_size - 1") + if self.head_offsets.shape != (heads,): + raise ValueError("Head offsets must match head vocabulary sizes") + for tensor in (self.multipliers, self.head_vocab_sizes, self.head_offsets): + if not tensor.is_cuda or not tensor.is_contiguous() or tensor.dtype != torch.int64: + raise ValueError("Hash buffers must be contiguous CUDA int64 tensors") + if tensor.device != self.multipliers.device: + raise ValueError("Hash buffers must share a device") + + sizes = self.head_vocab_sizes.tolist() + offsets = self.head_offsets.tolist() + previous_end = 0 + for size, offset in zip(sizes, offsets): + if size <= 0 or offset < previous_end or offset + size > torch.iinfo(torch.int64).max: + raise ValueError("Head vocabularies must have positive sizes and non-overlapping ordered offsets") + previous_end = offset + size + if not 0 <= self.eos_token_id <= torch.iinfo(torch.int64).max: + raise ValueError("EOS token ID must be a nonnegative int64 value") + + @property + def context_len(self) -> int: + return self.multipliers.numel() - 1 + + @property + def heads(self) -> int: + return self.head_vocab_sizes.numel() + + @classmethod + def from_checkpoint(cls, weights: Mapping[str, torch.Tensor], prefix: str, eos_token_id: int) -> "NGramHasher": + return cls( + weights[f"{prefix}.layer_multipliers"], + weights[f"{prefix}.ngram_heads_vocab_sizes"], + weights[f"{prefix}.ngram_heads_offsets"], + eos_token_id, + ) + + def empty_history(self, requests: int) -> torch.Tensor: + return torch.full( + (requests, self.context_len), self.eos_token_id, dtype=torch.int64, device=self.multipliers.device + ) + + def _validate(self, tokens: torch.Tensor, starts: torch.Tensor, history: torch.Tensor): + if tokens.ndim != 1 or tokens.dtype not in (torch.int32, torch.int64): + raise ValueError("Tokens must be a packed int32 or int64 vector") + if starts.ndim != 1 or starts.numel() == 0 or starts.dtype != torch.int32: + raise ValueError("Starts must be an int32 cumulative-length vector of length requests + 1") + requests = starts.numel() - 1 + if history.shape != (requests, self.context_len) or history.dtype != torch.int64: + raise ValueError("History must have shape (requests, ngram_size - 1) and dtype int64") + if requests == 0 and tokens.numel(): + raise ValueError("Nonempty tokens require at least one request") + for tensor in (tokens, starts, history): + if tensor.device != self.multipliers.device or not tensor.is_contiguous(): + raise ValueError("Inputs must be contiguous and on the hash buffers' device") + + @torch.no_grad() + def hash(self, tokens: torch.Tensor, starts: torch.Tensor, history: torch.Tensor): + self._validate(tokens, starts, history) + requests = starts.numel() - 1 + indices = torch.empty((tokens.numel(), self.heads), dtype=torch.int64, device=tokens.device) + if tokens.numel(): + _hash[(triton.cdiv(tokens.numel(), 64),)]( + tokens, + starts, + history, + self.multipliers, + self.head_vocab_sizes, + self.head_offsets, + indices, + tokens.numel(), + requests, + requests.bit_length(), + self.context_len, + self.heads, + self.heads // self.context_len, + self.eos_token_id, + 64, + triton.next_power_of_2(self.heads), + ) + return indices, self.history_after(tokens, starts, history) + + @torch.no_grad() + def history_after( + self, + tokens: torch.Tensor, + starts: torch.Tensor, + history: torch.Tensor, + accepted_lengths: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + self._validate(tokens, starts, history) + requests = starts.numel() - 1 + if accepted_lengths is not None: + if accepted_lengths.shape != (requests,) or accepted_lengths.dtype != torch.int32: + raise ValueError("Accepted lengths must be an int32 vector with one count per request") + if accepted_lengths.device != tokens.device or not accepted_lengths.is_contiguous(): + raise ValueError("Accepted lengths must be contiguous and on the input device") + output = torch.empty_like(history) + if requests: + _history[(requests,)]( + tokens, + starts, + history, + accepted_lengths, + output, + self.context_len, + accepted_lengths is not None, + self.eos_token_id, + triton.next_power_of_2(self.context_len), + num_warps=1, + ) + return output diff --git a/lightllm/common/ngram_embedding/kernels.py b/lightllm/common/ngram_embedding/kernels.py new file mode 100644 index 0000000000..f20f28c5b5 --- /dev/null +++ b/lightllm/common/ngram_embedding/kernels.py @@ -0,0 +1,94 @@ +import triton +import triton.language as tl + + +@triton.jit +def request_for_token(token, Starts, valid, REQUESTS: tl.constexpr, SEARCH_STEPS: tl.constexpr): + lo = tl.full(token.shape, 0, tl.int32) + hi = tl.full(token.shape, REQUESTS, tl.int32) + for _ in tl.static_range(SEARCH_STEPS): + mid = (lo + hi) // 2 + start = tl.load(Starts + mid, (mid < REQUESTS) & valid, other=2147483647) + right = start <= token + active = lo < hi + lo = tl.where(active & right, mid + 1, lo) + hi = tl.where(active & ~right, mid, hi) + request = tl.maximum(lo - 1, 0) + start = tl.load(Starts + request, valid, other=0) + return request, start + + +@triton.jit +def _hash( + Tokens, + Starts, + History, + Multipliers, + Sizes, + Offsets, + Output, + TOKENS: tl.constexpr, + REQUESTS: tl.constexpr, + SEARCH_STEPS: tl.constexpr, + CONTEXT: tl.constexpr, + HEADS: tl.constexpr, + HEADS_PER_NGRAM: tl.constexpr, + EOS: tl.constexpr, + BLOCK_T: tl.constexpr, + BLOCK_H: tl.constexpr, +): + token = tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T) + request, start = request_for_token(token, Starts, token < TOKENS, REQUESTS, SEARCH_STEPS) + position = token - start + head = tl.arange(0, BLOCK_H) + ngram_order = head // HEADS_PER_NGRAM + 2 + current = tl.load(Tokens + token, token < TOKENS, other=EOS).to(tl.int64) + multiplier = tl.load(Multipliers) + mixed = tl.broadcast_to((current * multiplier)[:, None], (BLOCK_T, BLOCK_H)) + alive = tl.full((BLOCK_T,), True, tl.int1) + for shift in tl.static_range(1, CONTEXT + 1): + in_chunk = position >= shift + source = tl.load(Tokens + token - shift, (token < TOKENS) & in_chunk, other=EOS).to(tl.int64) + from_history = tl.load( + History + request * CONTEXT + CONTEXT + position - shift, + (token < TOKENS) & ~in_chunk, + other=EOS, + ) + source = tl.where(in_chunk, source, from_history) + shifted = tl.where(alive, source, EOS) + multiplier = tl.load(Multipliers + shift) + mixed = mixed ^ tl.where(ngram_order[None, :] > shift, (shifted * multiplier)[:, None], 0) + + alive = alive & (source != EOS) + sizes = tl.load(Sizes + head, head < HEADS, other=1) + offsets = tl.load(Offsets + head, head < HEADS, other=0) + remainder = mixed % sizes[None, :] + remainder = tl.where(remainder < 0, remainder + sizes[None, :], remainder) + tl.store( + Output + token[:, None] * HEADS + head[None, :], + remainder + offsets[None, :], + (token[:, None] < TOKENS) & (head[None, :] < HEADS), + ) + + +@triton.jit +def _history( + Tokens, + Starts, + Previous, + Accepted, + Output, + CONTEXT: tl.constexpr, + USE_ACCEPTED: tl.constexpr, + EOS: tl.constexpr, + BLOCK: tl.constexpr, +): + request = tl.program_id(0) + col = tl.arange(0, BLOCK) + start = tl.load(Starts + request) + end = tl.load(Starts + request + 1) + count = tl.load(Accepted + request) if USE_ACCEPTED else end - start + source = count - CONTEXT + col + old = tl.load(Previous + request * CONTEXT + count + col, (col < CONTEXT) & (source < 0), other=EOS) + new = tl.load(Tokens + start + source, (col < CONTEXT) & (source >= 0), other=EOS).to(tl.int64) + tl.store(Output + request * CONTEXT + col, tl.where(source >= 0, new, old), col < CONTEXT) diff --git a/lightllm/common/ngram_embedding/ple.py b/lightllm/common/ngram_embedding/ple.py new file mode 100644 index 0000000000..042002cb8c --- /dev/null +++ b/lightllm/common/ngram_embedding/ple.py @@ -0,0 +1,198 @@ +from dataclasses import dataclass +from typing import Mapping, Optional + +import torch +import torch.nn.functional as F +import triton + +from lightllm.common.gated_residual.kernels import grouped_norm +from .ple_kernels import _gate, _conv, _conv_history + + +@dataclass(frozen=True) +class PLELayer: + key_weight: torch.Tensor + value_weight: torch.Tensor + norm_key: torch.Tensor + norm_query: torch.Tensor + norm_conv: torch.Tensor + conv_weight: torch.Tensor + branches: int + dilation: int + eps: float = 1e-6 + + def __post_init__(self): + width = self.norm_key.numel() + if self.branches < 1 or not width or width % self.branches or self.dilation < 1 or self.eps <= 0: + raise ValueError("PLE requires positive dimensions, dilation and epsilon") + if self.key_weight.ndim != 2 or self.key_weight.shape[1] == 0: + raise ValueError("PLE key projection must be a nonempty matrix") + if self.key_weight.shape[0] != width or self.value_weight.shape != (self.hidden_size, self.key_weight.shape[1]): + raise ValueError("PLE projection dimensions do not match residual/embedding widths") + if any(weight.shape != (width,) for weight in (self.norm_key, self.norm_query, self.norm_conv)): + raise ValueError("PLE normalization vectors must match the residual width") + if self.conv_weight.ndim != 2 or self.conv_weight.shape[0] != width or not self.conv_weight.shape[1]: + raise ValueError("PLE convolution weight must have shape (residual width, kernel size)") + for weight in ( + self.key_weight, + self.value_weight, + self.norm_key, + self.norm_query, + self.norm_conv, + self.conv_weight, + ): + if not weight.is_cuda or not weight.is_contiguous(): + raise ValueError("PLE weights must be contiguous CUDA tensors") + if weight.device != self.key_weight.device or weight.dtype != self.key_weight.dtype: + raise ValueError("PLE weights must share a device and dtype") + if self.key_weight.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError("PLE weights require a floating point dtype") + + @property + def hidden_size(self) -> int: + return self.norm_key.numel() // self.branches + + @property + def context_len(self) -> int: + return (self.conv_weight.shape[1] - 1) * self.dilation + + @classmethod + def from_checkpoint( + cls, weights: Mapping[str, torch.Tensor], prefix: str, branches: int, dilation: int, eps: float = 1e-6 + ) -> "PLELayer": + conv = weights[f"{prefix}.conv1d.weight"] + if conv.ndim != 3 or conv.shape[1] != 1: + raise ValueError("Checkpoint PLE convolution must be depthwise with shape (width, 1, kernel)") + return cls( + weights[f"{prefix}.key_proj.weight"], + weights[f"{prefix}.value_proj.weight"], + weights[f"{prefix}.norm_key.weight"], + weights[f"{prefix}.norm_query.weight"], + weights[f"{prefix}.norm_conv.weight"], + conv.squeeze(1), + branches, + dilation, + eps, + ) + + def empty_history(self, requests: int) -> torch.Tensor: + return torch.zeros( + (requests, self.context_len, self.norm_key.numel()), + dtype=self.key_weight.dtype, + device=self.key_weight.device, + ) + + def _check_tensor(self, tensor: torch.Tensor): + if ( + tensor.device != self.key_weight.device + or tensor.dtype != self.key_weight.dtype + or not tensor.is_contiguous() + ): + raise ValueError("PLE inputs must be contiguous and share the weights' device and dtype") + + def _validate_state(self, values: torch.Tensor, starts: torch.Tensor, history: torch.Tensor): + width = self.norm_key.numel() + self._check_tensor(values) + self._check_tensor(history) + if values.ndim != 2 or values.shape[1] != width: + raise ValueError("PLE values must have shape (packed tokens, residual width)") + if starts.ndim != 1 or starts.numel() == 0 or starts.dtype != torch.int32: + raise ValueError("PLE starts must be an int32 cumulative-length vector") + if starts.device != values.device or not starts.is_contiguous(): + raise ValueError("PLE starts must be contiguous and on the input device") + requests = starts.numel() - 1 + if history.shape != (requests, self.context_len, width): + raise ValueError("PLE history must have shape (requests, context length, residual width)") + if requests == 0 and values.shape[0]: + raise ValueError("Nonempty PLE input requires at least one request") + + @torch.no_grad() + def forward( + self, + residual: torch.Tensor, + embeddings: torch.Tensor, + starts: torch.Tensor, + history: torch.Tensor, + mask: Optional[torch.Tensor] = None, + ): + self._validate_state(residual, starts, history) + self._check_tensor(embeddings) + tokens = residual.shape[0] + if embeddings.shape != (tokens, self.key_weight.shape[1]): + raise ValueError("PLE embeddings must have shape (packed tokens, embedding width)") + if mask is not None: + if ( + mask.shape != (tokens,) + or mask.dtype != torch.bool + or mask.device != residual.device + or not mask.is_contiguous() + ): + raise ValueError("PLE mask must be a contiguous boolean vector on the input device") + key = grouped_norm(F.linear(embeddings, self.key_weight), self.norm_key, self.branches, self.eps) + query = grouped_norm(residual, self.norm_query, self.branches, self.eps) + value = F.linear(embeddings, self.value_weight) + gated = torch.empty_like(residual) + if tokens: + _gate[(tokens, self.branches)]( + key, + query, + value, + mask, + gated, + self.hidden_size, + self.branches, + mask is not None, + triton.next_power_of_2(self.hidden_size), + enable_fp_fusion=False, + ) + + conv_inputs = grouped_norm(gated, self.norm_conv, self.branches, self.eps) + output = torch.empty_like(residual) + if tokens: + requests = starts.numel() - 1 + _conv[(tokens, triton.cdiv(residual.shape[1], 256))]( + conv_inputs, + gated, + self.conv_weight, + starts, + history, + output, + residual.shape[1], + self.conv_weight.shape[1], + self.dilation, + requests, + requests.bit_length(), + 256, + enable_fp_fusion=False, + ) + return output, conv_inputs + + @torch.no_grad() + def history_after( + self, + conv_inputs: torch.Tensor, + starts: torch.Tensor, + history: torch.Tensor, + accepted_lengths: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + self._validate_state(conv_inputs, starts, history) + requests = starts.numel() - 1 + if accepted_lengths is not None: + if accepted_lengths.shape != (requests,) or accepted_lengths.dtype != torch.int32: + raise ValueError("PLE accepted lengths must be an int32 vector with one count per request") + if accepted_lengths.device != conv_inputs.device or not accepted_lengths.is_contiguous(): + raise ValueError("PLE accepted lengths must be contiguous and on the input device") + result = torch.empty_like(history) + if requests and self.context_len: + _conv_history[(requests, triton.cdiv(self.context_len * self.norm_key.numel(), 256))]( + conv_inputs, + starts, + history, + accepted_lengths, + result, + self.norm_key.numel(), + self.context_len, + accepted_lengths is not None, + 256, + ) + return result diff --git a/lightllm/common/ngram_embedding/ple_kernels.py b/lightllm/common/ngram_embedding/ple_kernels.py new file mode 100644 index 0000000000..79f7e13bac --- /dev/null +++ b/lightllm/common/ngram_embedding/ple_kernels.py @@ -0,0 +1,100 @@ +import triton +import triton.language as tl + +from .kernels import request_for_token + + +@triton.jit +def _gate(Key, Query, Value, Mask, Output, D: tl.constexpr, C: tl.constexpr, HAS_MASK: tl.constexpr, B: tl.constexpr): + token = tl.program_id(0).to(tl.int64) + branch = tl.program_id(1) + col = tl.arange(0, B) + offset = token * C * D + branch * D + col + key = tl.load(Key + offset, col < D, other=0).to(tl.float32) + query = tl.load(Query + offset, col < D, other=0).to(tl.float32) + + product = (key * query).to(Key.dtype.element_ty).to(tl.float32) + dot = tl.sum(product, 0).to(Key.dtype.element_ty).to(tl.float32) + dot = (dot / tl.sqrt(D * 1.0)).to(Key.dtype.element_ty).to(tl.float32) + magnitude = tl.maximum(tl.abs(dot), 1e-6).to(Key.dtype.element_ty).to(tl.float32) + root = tl.sqrt(magnitude).to(Key.dtype.element_ty).to(tl.float32) + signed = tl.where(dot > 0, root, tl.where(dot < 0, -root, 0.0)) + gate = tl.sigmoid(signed).to(Key.dtype.element_ty).to(tl.float32) + value = tl.load(Value + token * D + col, col < D, other=0).to(tl.float32) + result = value * gate + if HAS_MASK: + keep = tl.load(Mask + token) + result = tl.where(keep, result, 0.0) + tl.store(Output + offset, result, col < D) + + +@triton.jit +def _conv( + Input, + Gated, + Weight, + Starts, + History, + Output, + WIDTH: tl.constexpr, + KERNEL: tl.constexpr, + DILATION: tl.constexpr, + REQUESTS: tl.constexpr, + SEARCH_STEPS: tl.constexpr, + BLOCK: tl.constexpr, +): + token = tl.program_id(0).to(tl.int64) + col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + request, start = request_for_token(token, Starts, True, REQUESTS, SEARCH_STEPS) + context: tl.constexpr = (KERNEL - 1) * DILATION + position = token - start + acc = tl.full((BLOCK,), 0, tl.float32) + for tap in tl.static_range(KERNEL): + source = position - context + tap * DILATION + current = tl.load(Input + (start + source) * WIDTH + col, (col < WIDTH) & (source >= 0), other=0) + previous = tl.load( + History + (request.to(tl.int64) * context + context + source) * WIDTH + col, + (col < WIDTH) & (source < 0), + other=0, + ) + value = tl.where(source >= 0, current, previous).to(tl.float32) + weight = tl.load(Weight + col * KERNEL + tap, col < WIDTH, other=0).to(tl.float32) + acc = acc + value * weight + + conv = acc.to(Input.dtype.element_ty).to(tl.float32) + activated = (conv * tl.sigmoid(conv)).to(Input.dtype.element_ty).to(tl.float32) + gated = tl.load(Gated + token * WIDTH + col, col < WIDTH, other=0).to(tl.float32) + tl.store(Output + token * WIDTH + col, gated + activated, col < WIDTH) + + +@triton.jit +def _conv_history( + Input, + Starts, + Previous, + Accepted, + Output, + WIDTH: tl.constexpr, + CONTEXT: tl.constexpr, + USE_ACCEPTED: tl.constexpr, + BLOCK: tl.constexpr, +): + request = tl.program_id(0).to(tl.int64) + flat = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + time = flat // WIDTH + col = flat % WIDTH + start = tl.load(Starts + request) + end = tl.load(Starts + request + 1) + count = tl.load(Accepted + request) if USE_ACCEPTED else end - start + source = count - CONTEXT + time + old = tl.load( + Previous + (request * CONTEXT + count + time) * WIDTH + col, + (time < CONTEXT) & (source < 0), + other=0, + ) + new = tl.load( + Input + (start + source.to(tl.int64)) * WIDTH + col, + (time < CONTEXT) & (source >= 0), + other=0, + ) + tl.store(Output + request * CONTEXT * WIDTH + flat, tl.where(source >= 0, new, old), time < CONTEXT) diff --git a/lightllm/common/qsa/__init__.py b/lightllm/common/qsa/__init__.py new file mode 100644 index 0000000000..2ae28399f5 --- /dev/null +++ b/lightllm/common/qsa/__init__.py @@ -0,0 +1 @@ +pass diff --git a/lightllm/common/qsa/attention.py b/lightllm/common/qsa/attention.py new file mode 100644 index 0000000000..91815a0d63 --- /dev/null +++ b/lightllm/common/qsa/attention.py @@ -0,0 +1,120 @@ +from typing import Optional + +import torch +import triton + +from .attention_kernels import _attention, _reduce + + +@torch.no_grad() +def indexed_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + indices: torch.Tensor, + scale: Optional[float] = None, + splits: Optional[int] = None, + selection_lengths: Optional[torch.Tensor] = None, + short_rows_direct: bool = False, + skip_empty_splits: bool = False, + block_size: int = 64, + num_stages: Optional[int] = None, +) -> torch.Tensor: + if query.ndim != 3 or key.ndim != 3 or value.ndim != 3 or indices.ndim != 2: + raise ValueError("Indexed attention requires rank-three Q/K/V and rank-two indices") + rows, q_heads, dim = query.shape + if not q_heads or not dim or not key.shape[1] or q_heads % key.shape[1]: + raise ValueError("Indexed attention query heads must be divisible by KV heads") + if key.shape[0:2] != value.shape[0:2] or key.shape[2] != dim or not value.shape[2]: + raise ValueError("Indexed attention key/value dimensions do not match") + for tensor in (query, key, value): + if not tensor.is_cuda or tensor.dtype != query.dtype or tensor.device != query.device: + raise ValueError("Indexed attention Q/K/V must be CUDA tensors with matching dtype/device") + if not query.is_contiguous(): + raise ValueError("Indexed attention query must be contiguous") + if query.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError("Indexed attention requires floating point inputs") + if ( + indices.shape[0] != rows + or indices.dtype != torch.int32 + or indices.device != query.device + or not indices.is_contiguous() + ): + raise ValueError("Indexed attention indices must be a contiguous int32 matrix on the query device") + selected = indices.shape[1] + if selection_lengths is not None and ( + selection_lengths.shape != (rows,) + or selection_lengths.dtype != torch.int32 + or selection_lengths.device != query.device + or not selection_lengths.is_contiguous() + ): + raise ValueError("Selection lengths must be a contiguous int32 vector on the query device") + kv_heads, value_dim = value.shape[1:] + output = torch.empty((rows, q_heads, value_dim), dtype=query.dtype, device=query.device) + if not rows: + return output + if not selected: + return output.zero_() + if block_size < 16 or block_size & (block_size - 1): + raise ValueError("Indexed attention block size must be a power of two >= 16") + if num_stages is not None and num_stages < 1: + raise ValueError("Indexed attention num_stages must be positive") + if splits is None: + splits = min(32, max(1, triton.cdiv(block_size, rows * kv_heads)), triton.cdiv(selected, block_size)) + if splits < 1: + raise ValueError("Indexed attention split count must be positive") + partial = ( + torch.empty((rows, q_heads, splits, value_dim), dtype=torch.float32, device=query.device) + if splits > 1 + else None + ) + lse = torch.empty((rows, q_heads, splits), dtype=torch.float32, device=query.device) if splits > 1 else None + launch_kwargs = {} + if num_stages is not None: + launch_kwargs["num_stages"] = num_stages + _attention[(rows, kv_heads, splits)]( + query, + key, + value, + indices, + selection_lengths, + partial, + lse, + output, + q_heads, + kv_heads, + *key.stride(), + *value.stride(), + q_heads // kv_heads, + dim, + value_dim, + selected, + splits, + dim ** -0.5 if scale is None else scale, + max(16, triton.next_power_of_2(q_heads // kv_heads)), + max(16, triton.next_power_of_2(dim)), + max(16, triton.next_power_of_2(value_dim)), + block_size, + short_rows_direct, + skip_empty_splits, + "ieee" if query.dtype == torch.float32 else "tf32", + **launch_kwargs, + ) + if splits > 1: + _reduce[(rows, q_heads)]( + selection_lengths, + partial, + lse, + output, + q_heads, + value_dim, + splits, + triton.next_power_of_2(splits), + triton.next_power_of_2(value_dim), + selected, + block_size, + short_rows_direct, + skip_empty_splits, + **launch_kwargs, + ) + return output diff --git a/lightllm/common/qsa/attention_kernels.py b/lightllm/common/qsa/attention_kernels.py new file mode 100644 index 0000000000..7ce95104b0 --- /dev/null +++ b/lightllm/common/qsa/attention_kernels.py @@ -0,0 +1,137 @@ +import triton +import triton.language as tl + + +@triton.jit +def _attention( + Q, + K, + V, + Indices, + Lengths, + Partial, + LSE, + Output, + Q_HEADS: tl.constexpr, + KV_HEADS: tl.constexpr, + KS: tl.constexpr, + KH: tl.constexpr, + KD: tl.constexpr, + VS: tl.constexpr, + VH: tl.constexpr, + VD: tl.constexpr, + GROUP: tl.constexpr, + D: tl.constexpr, + DV: tl.constexpr, + SELECTED: tl.constexpr, + SPLITS: tl.constexpr, + SCALE: tl.constexpr, + BG: tl.constexpr, + BD: tl.constexpr, + BV: tl.constexpr, + BN: tl.constexpr, + SHORT_DIRECT: tl.constexpr, + SKIP_EMPTY: tl.constexpr, + INPUT_PRECISION: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + kv_head = tl.program_id(1) + split = tl.program_id(2) + limit = SELECTED + if Lengths is not None: + limit = tl.minimum(SELECTED, tl.maximum(0, tl.load(Lengths + row))) + if (not SHORT_DIRECT or limit > BN or split == 0) and ( + not SKIP_EMPTY or SPLITS == 1 or split * BN < limit or (SHORT_DIRECT and limit <= BN and split == 0) + ): + group = tl.arange(0, BG) + head = kv_head * GROUP + group + dim = tl.arange(0, BD) + value_dim = tl.arange(0, BV) + q = tl.load( + Q + (row * Q_HEADS + head[:, None]) * D + dim[None, :], + (group[:, None] < GROUP) & (dim[None, :] < D), + other=0, + ) + maximum = tl.full((BG,), -float("inf"), tl.float32) + denominator = tl.full((BG,), 0, tl.float32) + acc = tl.full((BG, BV), 0, tl.float32) + for start in tl.range(split * BN, limit, SPLITS * BN): + selected = start + tl.arange(0, BN) + slot = tl.load(Indices + row * SELECTED + selected, selected < limit, other=-1).to(tl.int64) + valid = (selected < limit) & (slot >= 0) + key = tl.load( + K + slot[None, :] * KS + kv_head * KH + dim[:, None] * KD, valid[None, :] & (dim[:, None] < D), other=0 + ) + score = tl.dot(q, key, input_precision=INPUT_PRECISION) * SCALE + score = tl.where(valid[None, :], score, -float("inf")) + new_max = tl.maximum(maximum, tl.max(score, 1)) + + safe_max = tl.where(new_max == -float("inf"), 0.0, new_max) + alpha = tl.exp(maximum - safe_max) + prob = tl.exp(score - safe_max[:, None]) + value = tl.load( + V + slot[:, None] * VS + kv_head * VH + value_dim[None, :] * VD, + valid[:, None] & (value_dim[None, :] < DV), + other=0, + ) + acc = acc * alpha[:, None] + acc = tl.dot(prob.to(V.dtype.element_ty), value, acc, input_precision=INPUT_PRECISION) + denominator = denominator * alpha + tl.sum(prob, 1) + maximum = new_max + safe_denominator = tl.where(denominator > 0, denominator, 1.0) + result = acc / safe_denominator[:, None] + if SPLITS == 1 or (SHORT_DIRECT and limit <= BN): + tl.store( + Output + (row * Q_HEADS + head[:, None]) * DV + value_dim[None, :], + result, + (group[:, None] < GROUP) & (value_dim[None, :] < DV), + ) + else: + tl.store( + Partial + ((row * Q_HEADS + head[:, None]) * SPLITS + split) * DV + value_dim[None, :], + result, + (group[:, None] < GROUP) & (value_dim[None, :] < DV), + ) + lse = tl.where(denominator > 0, maximum + tl.log(safe_denominator), -float("inf")) + tl.store(LSE + (row * Q_HEADS + head) * SPLITS + split, lse, group < GROUP) + + +@triton.jit +def _reduce( + Lengths, + Partial, + LSE, + Output, + HEADS: tl.constexpr, + DV: tl.constexpr, + SPLITS: tl.constexpr, + BS: tl.constexpr, + BV: tl.constexpr, + SELECTED: tl.constexpr, + BN: tl.constexpr, + SHORT_DIRECT: tl.constexpr, + SKIP_EMPTY: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + limit = SELECTED + if Lengths is not None: + limit = tl.minimum(SELECTED, tl.maximum(0, tl.load(Lengths + row))) + if not SHORT_DIRECT or limit > BN: + head = tl.program_id(1) + split = tl.arange(0, BS) + dim = tl.arange(0, BV) + valid_split = split < SPLITS + if SKIP_EMPTY: + valid_split = valid_split & (split * BN < limit) + lse = tl.load(LSE + (row * HEADS + head) * SPLITS + split, valid_split, other=-float("inf")) + maximum = tl.max(lse, 0) + maximum = tl.where(maximum == -float("inf"), 0.0, maximum) + weights = tl.exp(lse - maximum) + partial = tl.load( + Partial + ((row * HEADS + head) * SPLITS + split[:, None]) * DV + dim[None, :], + valid_split[:, None] & (dim[None, :] < DV), + other=0, + ) + total = tl.sum(weights, 0) + result = tl.sum(partial * weights[:, None], 0) / tl.where(total > 0, total, 1.0) + tl.store(Output + (row * HEADS + head) * DV + dim, result, dim < DV) diff --git a/lightllm/common/qsa/indexer.py b/lightllm/common/qsa/indexer.py new file mode 100644 index 0000000000..eda8a9158b --- /dev/null +++ b/lightllm/common/qsa/indexer.py @@ -0,0 +1,279 @@ +from dataclasses import dataclass +from typing import Mapping, Optional + +import torch +import torch.nn.functional as F +import triton + +from lightllm.common.gated_residual.kernels import grouped_norm +from .indexer_kernels import _rotate, _pool, _scores, _expand +from .topk import cuda_topk_indices, cuda_topk_supported, flashinfer_topk_indices, flashinfer_topk_supported + + +@dataclass(frozen=True) +class QSAIndexer: + projection: torch.Tensor + query_norm: torch.Tensor + key_norm: torch.Tensor + query_heads: int = 4 + token_budget: int = 2048 + compress_ratio: int = 4 + eps: float = 1e-6 + use_cuda_topk: bool = False + + def __post_init__(self): + dim = self.key_norm.numel() + if not dim or self.query_heads < 1 or self.compress_ratio < 1 or self.token_budget < 1 or self.eps <= 0: + raise ValueError("QSA dimensions, budget, compression ratio and epsilon must be positive") + if self.token_budget % self.compress_ratio: + raise ValueError("QSA token budget must be divisible by compression ratio") + if self.key_norm.shape != (dim,) or self.query_norm.shape != (dim,): + raise ValueError("QSA normalization weights must be matching vectors") + if ( + self.projection.ndim != 2 + or self.projection.shape[0] != (self.query_heads + 1) * dim + or not self.projection.shape[1] + ): + raise ValueError("QSA projection must contain query heads and one raw key head") + for weight in (self.projection, self.query_norm, self.key_norm): + self._activation(weight) + + @property + def head_dim(self) -> int: + return self.key_norm.numel() + + @classmethod + def from_checkpoint(cls, weights: Mapping[str, torch.Tensor], prefix: str, **kwargs) -> "QSAIndexer": + return cls( + weights[f"{prefix}.index_qk_proj.weight"], + weights[f"{prefix}.q_layernorm.weight"], + weights[f"{prefix}.k_layernorm.weight"], + **kwargs, + ) + + def _activation(self, x: torch.Tensor): + if ( + not x.is_cuda + or not x.is_contiguous() + or x.device != self.projection.device + or x.dtype != self.projection.dtype + ): + raise ValueError("QSA activations/weights must be contiguous CUDA tensors on the same device and dtype") + if x.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError("QSA requires floating point activations") + + def _integer(self, x: torch.Tensor): + if x.dtype != torch.int32 or x.device != self.projection.device or not x.is_contiguous(): + raise ValueError("QSA metadata must be contiguous int32 tensors on the input device") + + def _rotary(self, cos: torch.Tensor, sin: torch.Tensor, rows: int) -> int: + self._activation(cos) + self._activation(sin) + if cos.ndim != 2 or sin.shape != cos.shape or cos.shape[0] != rows: + raise ValueError("QSA rotary tensors must have matching shapes (tokens, rotary dimensions)") + if cos.shape[1] > self.head_dim or cos.shape[1] % 2: + raise ValueError("QSA rotary dimensions must be even and fit the index head") + return cos.shape[1] + + @torch.no_grad() + def project(self, hidden: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): + self._activation(hidden) + if hidden.ndim != 2 or hidden.shape[1] != self.projection.shape[1]: + raise ValueError("QSA hidden input must have shape (tokens, hidden size)") + rotary_dim = self._rotary(cos, sin, hidden.shape[0]) + projected = F.linear(hidden, self.projection) + query = projected[:, : self.query_heads * self.head_dim].contiguous().view(-1, self.head_dim) + query = grouped_norm(query, self.query_norm, 1, self.eps).view(hidden.shape[0], self.query_heads, self.head_dim) + output = torch.empty_like(query) + if hidden.shape[0]: + _rotate[(hidden.shape[0], self.query_heads)]( + query, + cos, + sin, + output, + self.query_heads, + self.head_dim, + rotary_dim, + triton.next_power_of_2(self.head_dim), + enable_fp_fusion=False, + ) + # Keep the raw-key projection as a view. The state writer accepts a + # strided source, so making this slice contiguous only adds a full + # token-by-head-dimension copy on every QSA layer. + return output, projected[:, self.query_heads * self.head_dim :] + + def empty_cache(self, requests: int, max_tokens: int) -> torch.Tensor: + return torch.empty( + (requests, max_tokens // self.compress_ratio, self.head_dim), + dtype=self.projection.dtype, + device=self.projection.device, + ) + + def empty_physical_cache(self, slots: int) -> torch.Tensor: + if slots < 0: + raise ValueError("QSA physical cache size cannot be negative") + return torch.empty((slots, self.head_dim), dtype=self.projection.dtype, device=self.projection.device) + + @torch.no_grad() + def pool( + self, + raw_keys: torch.Tensor, + indices: torch.Tensor, + lengths: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + cache: torch.Tensor, + first_blocks: torch.Tensor, + max_new_blocks: int, + request_rows: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + self._activation(raw_keys) + self._activation(cache) + for tensor in (indices, lengths, first_blocks): + self._integer(tensor) + if raw_keys.ndim != 2 or raw_keys.shape[1] != self.head_dim or indices.ndim != 2: + raise ValueError("QSA raw keys and slot indices must be matrices") + requests, capacity = indices.shape + active = requests + if request_rows is not None: + self._integer(request_rows) + if request_rows.ndim != 1: + raise ValueError("QSA active request rows must be a vector") + active = request_rows.numel() + if lengths.shape != (active,) or first_blocks.shape != (active,): + raise ValueError("QSA lengths and first blocks must have one entry per request") + physical = cache.ndim == 2 + if physical: + if cache.shape != raw_keys.shape: + raise ValueError("QSA physical pooled cache must match the raw key namespace") + elif ( + cache.ndim != 3 + or cache.shape[0] != requests + or cache.shape[1] < capacity // self.compress_ratio + or cache.shape[2] != self.head_dim + ): + raise ValueError("QSA cache must hold all complete blocks in each request's slot table") + if max_new_blocks < 0: + raise ValueError("QSA maximum new block count cannot be negative") + rotary_dim = self._rotary(cos, sin, raw_keys.shape[0]) + if active and max_new_blocks: + _pool[(active, max_new_blocks)]( + raw_keys, + indices, + lengths, + first_blocks, + self.key_norm, + cos, + sin, + cache, + request_rows, + request_rows is not None, + physical, + capacity, + cache.shape[1], + self.head_dim, + rotary_dim, + self.compress_ratio, + self.eps, + triton.next_power_of_2(self.head_dim), + triton.next_power_of_2(self.compress_ratio), + enable_fp_fusion=False, + ) + return cache + + @torch.no_grad() + def select( + self, + query: torch.Tensor, + cache: torch.Tensor, + requests: torch.Tensor, + visible_lengths: torch.Tensor, + indices: torch.Tensor, + query_chunk_size: int = 64, + ) -> torch.Tensor: + self._activation(query) + self._activation(cache) + for tensor in (requests, visible_lengths, indices): + self._integer(tensor) + if query.ndim != 3 or query.shape[1:] != (self.query_heads, self.head_dim): + raise ValueError("QSA query shape must be (tokens, index query heads, index head dimension)") + tokens = query.shape[0] + if requests.shape != (tokens,) or visible_lengths.shape != (tokens,) or indices.ndim != 2: + raise ValueError("QSA requires one request ID and visible length per query") + physical = cache.ndim == 2 + if physical: + if cache.shape[1] != self.head_dim: + raise ValueError("QSA physical pooled cache has the wrong head dimension") + elif ( + cache.ndim != 3 + or cache.shape[0] != indices.shape[0] + or cache.shape[2] != self.head_dim + or cache.shape[1] < indices.shape[1] // self.compress_ratio + ): + raise ValueError("QSA pooled cache and physical slot table dimensions do not match") + if query_chunk_size < 1: + raise ValueError("QSA query chunk size must be positive") + width = self.token_budget + self.compress_ratio - 1 + output = torch.empty((tokens, width), dtype=torch.int32, device=query.device) + blocks = min(indices.shape[1], cache.shape[0]) // self.compress_ratio if physical else cache.shape[1] + block_budget = self.token_budget // self.compress_ratio + topk = min(blocks, block_budget) + cuda_topk = self.use_cuda_topk and cuda_topk_supported(query.device, blocks, topk) + flashinfer_topk = self.use_cuda_topk and not cuda_topk and flashinfer_topk_supported(query.device, blocks, topk) + max_score_elements = 16 * 1024 * 1024 + query_chunk_size = min(query_chunk_size, max(1, max_score_elements // max(blocks, 1))) + score_buffer = torch.empty((min(tokens, query_chunk_size), blocks), dtype=torch.float32, device=query.device) + group_size = 4 if physical and self.query_heads == 4 else 1 + score_block = 128 if group_size == 4 else 64 + for first in range(0, tokens, query_chunk_size): + count = min(query_chunk_size, tokens - first) + scores = score_buffer[:count] + if blocks: + _scores[(triton.cdiv(count, group_size), triton.cdiv(blocks, score_block))]( + query, + cache, + requests, + visible_lengths, + indices, + scores, + first, + self.query_heads, + self.head_dim, + blocks, + indices.shape[1], + physical, + self.compress_ratio, + max(16, triton.next_power_of_2(self.query_heads)), + max(16, triton.next_power_of_2(self.head_dim)), + score_block, + block_budget if cuda_topk or flashinfer_topk else 0, + "ieee" if query.dtype == torch.float32 else "tf32", + GROUP_SIZE=group_size, + TOKENS=count, + num_warps=4, + ) + if cuda_topk: + selected = cuda_topk_indices( + scores, topk, visible_lengths=visible_lengths, first=first, compress_ratio=self.compress_ratio + ) + elif flashinfer_topk: + selected = flashinfer_topk_indices( + scores, topk, visible_lengths=visible_lengths, first=first, compress_ratio=self.compress_ratio + ) + else: + selected = torch.topk(scores, topk, dim=-1, sorted=False).indices + _expand[(count,)]( + selected, + requests, + visible_lengths, + indices, + output, + first, + topk, + block_budget, + indices.shape[1], + self.compress_ratio, + width, + triton.next_power_of_2(width), + ) + return output diff --git a/lightllm/common/qsa/indexer_kernels.py b/lightllm/common/qsa/indexer_kernels.py new file mode 100644 index 0000000000..60a98de0dc --- /dev/null +++ b/lightllm/common/qsa/indexer_kernels.py @@ -0,0 +1,216 @@ +import triton +import triton.language as tl + + +@triton.jit +def _rotate(X, Cos, Sin, Output, HEADS: tl.constexpr, D: tl.constexpr, R: tl.constexpr, BLOCK: tl.constexpr): + row = tl.program_id(0).to(tl.int64) + head = tl.program_id(1) + col = tl.arange(0, BLOCK) + x = tl.load(X + (row * HEADS + head) * D + col, col < D, other=0) + if R > 0: + partner = tl.where(col < R // 2, col + R // 2, col - R // 2) + other = tl.load(X + (row * HEADS + head) * D + partner, col < R, other=0).to(tl.float32) + other = tl.where(col < R // 2, -other, other) + cos = tl.load(Cos + row * R + col, col < R, other=0).to(tl.float32) + sin = tl.load(Sin + row * R + col, col < R, other=0).to(tl.float32) + left = (x.to(tl.float32) * cos).to(X.dtype.element_ty).to(tl.float32) + right = (other * sin).to(X.dtype.element_ty).to(tl.float32) + x = tl.where(col < R, left + right, x.to(tl.float32)) + tl.store(Output + (row * HEADS + head) * D + col, x, col < D) + + +@triton.jit +def _pool( + Raw, + Indices, + Lengths, + First, + Norm, + Cos, + Sin, + Cache, + RequestRows, + MAPPED_ROWS: tl.constexpr, + PHYSICAL_CACHE: tl.constexpr, + CAPACITY: tl.constexpr, + BLOCKS: tl.constexpr, + D: tl.constexpr, + R: tl.constexpr, + RATIO: tl.constexpr, + EPS: tl.constexpr, + BD: tl.constexpr, + BR: tl.constexpr, +): + batch_row = tl.program_id(0) + request = tl.load(RequestRows + batch_row).to(tl.int64) if MAPPED_ROWS else batch_row.to(tl.int64) + block = tl.load(First + batch_row) + tl.program_id(1) + count = tl.load(Lengths + batch_row) + if (request >= 0) & (block < count // RATIO): + token = tl.arange(0, BR) + col = tl.arange(0, BD) + slots = tl.load(Indices + request * CAPACITY + block * RATIO + token, token < RATIO, other=0).to(tl.int64) + values = tl.load( + Raw + slots[:, None] * D + col[None, :], (token[:, None] < RATIO) & (col[None, :] < D), other=0 + ).to(tl.float32) + pooled = (tl.sum(values, 0) / RATIO).to(Raw.dtype.element_ty).to(tl.float32) + inv = tl.rsqrt(tl.sum(pooled * pooled, 0) / D + EPS) + weight = tl.load(Norm + col, col < D, other=0).to(tl.float32) + normalized = ((pooled * inv) * (1 + weight)).to(Raw.dtype.element_ty).to(tl.float32) + if R > 0: + slot = tl.load(Indices + request * CAPACITY + block * RATIO).to(tl.int64) + partner = tl.where(col < R // 2, col + R // 2, col - R // 2) + + partner = tl.minimum(tl.maximum(partner, 0), BD - 1) + other = tl.gather(normalized, partner, 0) + other = tl.where(col < R // 2, -other, other) + cos = tl.load(Cos + slot * R + col, col < R, other=0).to(tl.float32) + sin = tl.load(Sin + slot * R + col, col < R, other=0).to(tl.float32) + left = (normalized * cos).to(Raw.dtype.element_ty).to(tl.float32) + right = (other * sin).to(Raw.dtype.element_ty).to(tl.float32) + normalized = tl.where(col < R, left + right, normalized) + if PHYSICAL_CACHE: + destination = tl.load(Indices + request * CAPACITY + block * RATIO + RATIO - 1).to(tl.int64) + else: + destination = request * BLOCKS + block + tl.store(Cache + destination * D + col, normalized, col < D) + + +@triton.jit +def _scores( + Query, + Cache, + Requests, + Lengths, + Indices, + Output, + QUERY_START: tl.constexpr, + HEADS: tl.constexpr, + D: tl.constexpr, + BLOCKS: tl.constexpr, + CAPACITY: tl.constexpr, + PHYSICAL_CACHE: tl.constexpr, + RATIO: tl.constexpr, + BH: tl.constexpr, + BD: tl.constexpr, + BN: tl.constexpr, + MIN_BLOCKS: tl.constexpr = 0, + INPUT_PRECISION: tl.constexpr = "tf32", + GROUP_SIZE: tl.constexpr = 1, + TOKENS: tl.constexpr = 0, +): + tl.static_assert(GROUP_SIZE == 1 or (GROUP_SIZE == 4 and HEADS == 4 and PHYSICAL_CACHE)) + group = tl.program_id(0) + if GROUP_SIZE == 4: + rows = group * 4 + tl.arange(0, 4) + valid = rows < TOKENS + requests = tl.load(Requests + QUERY_START + rows, valid, 0).to(tl.int64) + lengths = tl.load(Lengths + QUERY_START + rows, valid, 0) // RATIO + shared = tl.min(tl.where(valid, requests, 0x7FFFFFFF)) == tl.max(tl.where(valid, requests, 0)) + else: + shared = tl.full((), False, tl.int1) + if GROUP_SIZE == 4 and shared: + # Four queries fill the MMA rows and share the same pooled keys. + request = tl.sum(tl.where(tl.arange(0, 4) == 0, requests, 0)) + maximum = tl.max(lengths) + block = tl.program_id(1) * BN + tl.arange(0, BN) + scores = tl.full((4, BN), -float("inf"), tl.float32) + if (tl.program_id(1) * BN < maximum) & (maximum >= MIN_BLOCKS): + head = tl.arange(0, 16) + col = tl.arange(0, BD) + query_rows = QUERY_START + group * 4 + head // HEADS + q = tl.load( + Query + (query_rows[:, None].to(tl.int64) * HEADS + (head % HEADS)[:, None]) * D + col[None, :], + (query_rows[:, None] < QUERY_START + TOKENS) & (col[None, :] < D), + 0, + ) + key_rows = tl.load( + Indices + request * CAPACITY + block * RATIO + RATIO - 1, + (block < maximum) & (block < BLOCKS), + 0, + ).to(tl.int64) + k = tl.load( + Cache + key_rows[None, :] * D + col[:, None], (block[None, :] < maximum) & (col[:, None] < D), 0 + ) + logits = tl.dot(q, k, input_precision=INPUT_PRECISION) + scores = tl.sum(tl.maximum(logits, 0).reshape(4, HEADS, BN), 1) * (D ** -0.5) + scores = tl.where(block[None, :] < lengths[:, None], scores, -float("inf")) + tl.store( + Output + rows[:, None].to(tl.int64) * BLOCKS + block[None, :], + scores, + valid[:, None] & (block[None, :] < BLOCKS) & ((MIN_BLOCKS == 0) | (lengths[:, None] >= MIN_BLOCKS)), + ) + else: + for offset in range(GROUP_SIZE): + row = group * GROUP_SIZE + offset + if GROUP_SIZE == 1 or row < TOKENS: + query_row = row + QUERY_START + request = tl.load(Requests + query_row).to(tl.int64) + visible = tl.load(Lengths + query_row) // RATIO + block = tl.program_id(1) * BN + tl.arange(0, BN) + + scores = tl.full((BN,), -float("inf"), tl.float32) + if (tl.program_id(1) * BN < visible) & (visible >= MIN_BLOCKS): + head = tl.arange(0, BH) + col = tl.arange(0, BD) + q = tl.load( + Query + (query_row.to(tl.int64) * HEADS + head[:, None]) * D + col[None, :], + (head[:, None] < HEADS) & (col[None, :] < D), + other=0, + ) + if PHYSICAL_CACHE: + key_rows = tl.load( + Indices + request * CAPACITY + block * RATIO + RATIO - 1, + (block < visible) & (block < BLOCKS), + other=0, + ).to(tl.int64) + else: + key_rows = request * BLOCKS + block + k = tl.load( + Cache + key_rows[None, :] * D + col[:, None], + (block[None, :] < visible) & (col[:, None] < D), + other=0, + ) + logits = tl.dot(q, k, input_precision=INPUT_PRECISION) + logits = tl.where(head[:, None] < HEADS, tl.maximum(logits, 0), 0) + scores = tl.sum(logits, 0) * (D ** -0.5) + scores = tl.where(block < visible, scores, -float("inf")) + tl.store( + Output + row.to(tl.int64) * BLOCKS + block, + scores, + (block < BLOCKS) & ((MIN_BLOCKS == 0) | (visible >= MIN_BLOCKS)), + ) + + +@triton.jit +def _expand( + Selected, + Requests, + Lengths, + Indices, + Output, + QUERY_START: tl.constexpr, + TOPK: tl.constexpr, + BLOCK_BUDGET: tl.constexpr, + CAPACITY: tl.constexpr, + RATIO: tl.constexpr, + WIDTH: tl.constexpr, + B: tl.constexpr, +): + row = tl.program_id(0) + query_row = row + QUERY_START + request = tl.load(Requests + query_row).to(tl.int64) + visible = tl.load(Lengths + query_row) + complete = visible // RATIO + chosen = tl.minimum(complete, BLOCK_BUDGET) + col = tl.arange(0, B) + rank = col // RATIO + block = tl.full((B,), 0, tl.int64) + if TOPK > 0: + block = tl.load(Selected + row.to(tl.int64) * TOPK + rank, (rank < chosen) & (col < WIDTH), other=0) + full = col < chosen * RATIO + tail_col = col - chosen * RATIO + tail = (tail_col >= 0) & (tail_col < visible % RATIO) + position = tl.where(full, block * RATIO + col % RATIO, complete * RATIO + tail_col) + slot = tl.load(Indices + request * CAPACITY + position, (col < WIDTH) & (full | tail), other=-1) + tl.store(Output + query_row.to(tl.int64) * WIDTH + col, slot, col < WIDTH) diff --git a/lightllm/common/qsa/topk.py b/lightllm/common/qsa/topk.py new file mode 100644 index 0000000000..c6c983d860 --- /dev/null +++ b/lightllm/common/qsa/topk.py @@ -0,0 +1,127 @@ +from functools import lru_cache +from inspect import signature + +import torch +import triton +import triton.language as tl +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +_CUDA_TOPK_RUNTIME = str(torch.__version__) == "2.9.1+cu128" and triton.__version__ == "3.5.1" + + +@lru_cache(maxsize=16) +def _supports_device(device): + supported = torch.cuda.get_device_capability(device) == (9, 0) and "H100" in torch.cuda.get_device_name(device) + if supported: + logger.info("QSA CUDA TopK selection enabled") + return supported + + +@triton.jit +def _key(value): + bits = value.to(tl.uint32, bitcast=True) + mask = tl.where((bits & 0x80000000) != 0, 0xFFFFFFFF, 0x80000000).to(tl.uint32) + return tl.where(value == value, bits ^ mask, 0xFFFFFFFF).to(tl.uint32) + + +@triton.jit +def _topk( + Scores, + Selected, + Lengths, + FIRST: tl.constexpr, + RATIO: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, + B: tl.constexpr, +): + tl.static_assert(N == 1027 and K == 512 and B == 2048) + row = tl.program_id(0).to(tl.int64) + if Lengths is None: + short = tl.full((), False, tl.int1) + else: + short = tl.load(Lengths + row + FIRST) // RATIO < K + if short: + short_col = tl.arange(0, K) + tl.store(Selected + row * K + short_col, short_col.to(tl.int64)) + else: + prefix_col = tl.arange(0, 1024) + prefix = tl.topk(_key(tl.load(Scores + row * N + prefix_col)), K) + merge_col = tl.arange(0, 4) + tail = tl.load(Scores + row * N + 1024 + merge_col, merge_col < 3, other=0.0) + tail = tl.sort(tl.where(merge_col < 3, _key(tail), 0).to(tl.uint32), descending=True) + + prefix_cutoffs = tl.gather(prefix, K - 1 - merge_col, 0) + tail_cutoffs = tl.gather(tail, tl.maximum(merge_col - 1, 0), 0) + tail_cutoffs = tl.where(merge_col == 0, 0xFFFFFFFF, tail_cutoffs).to(tl.uint32) + threshold = tl.max(tl.minimum(prefix_cutoffs, tail_cutoffs)) + col = tl.arange(0, B) + value = tl.load(Scores + row * N + col, col < N, other=-float("inf")) + key = _key(value) + greater = (col < N) & (key > threshold) + equal = (col < N) & (key == threshold) + greater_rank = tl.cumsum(greater.to(tl.int32)) - 1 + equal_rank = tl.cumsum(equal.to(tl.int32)) - 1 + rank = tl.where(greater, greater_rank, tl.sum(greater.to(tl.int32)) + equal_rank) + selected = greater | (equal & (rank < K)) + tl.store(Selected + row * K + rank, col.to(tl.int64), selected) + + +def cuda_topk_supported(device, blocks: int, k: int) -> bool: + return _CUDA_TOPK_RUNTIME and device.type == "cuda" and blocks == 1027 and k == 512 and _supports_device(device) + + +@lru_cache(maxsize=1) +def _flashinfer_topk(): + try: + from flashinfer import top_k_ragged_transform + except ImportError: + return None + if "dsa_graph_safe" not in signature(top_k_ragged_transform).parameters: + return None + return top_k_ragged_transform + + +@lru_cache(maxsize=16) +def flashinfer_topk_supported(device, blocks: int, k: int) -> bool: + return ( + device.type == "cuda" + and k == 512 + and blocks >= k + and torch.cuda.get_device_capability(device) == (9, 0) + and _flashinfer_topk() is not None + ) + + +def flashinfer_topk_indices(scores, k, *, visible_lengths, first=0, compress_ratio=4): + lengths = visible_lengths[first : first + scores.shape[0]] // compress_ratio + return _flashinfer_topk()(scores, torch.zeros_like(lengths), lengths, k, dsa_graph_safe=True) + + +def cuda_topk_indices(scores: torch.Tensor, k: int, *, visible_lengths=None, first=0, compress_ratio=4) -> torch.Tensor: + if not ( + scores.is_cuda + and scores.dtype == torch.float32 + and scores.ndim == 2 + and scores.is_contiguous() + and cuda_topk_supported(scores.device, scores.shape[1], k) + ): + return torch.topk(scores, k, dim=-1, sorted=False).indices + if visible_lengths is not None: + if ( + visible_lengths.dtype != torch.int32 + or visible_lengths.device != scores.device + or visible_lengths.ndim != 1 + or not visible_lengths.is_contiguous() + or first < 0 + or first + scores.shape[0] > visible_lengths.numel() + or compress_ratio < 1 + ): + raise ValueError("QSA TopK requires matching contiguous int32 visible lengths and valid bounds") + output = torch.empty((scores.shape[0], k), dtype=torch.int64, device=scores.device) + if scores.shape[0]: + _topk[(scores.shape[0],)](scores, output, visible_lengths, first, compress_ratio, 1027, 512, 2048, num_warps=16) + return output diff --git a/lightllm/common/quantization/aligned_projection.py b/lightllm/common/quantization/aligned_projection.py new file mode 100644 index 0000000000..06e1f9a287 --- /dev/null +++ b/lightllm/common/quantization/aligned_projection.py @@ -0,0 +1,51 @@ +import torch + +from lightllm.common.quantization.no_quant import NoQuantization +from lightllm.common.quantization.quantize_method import WeightPack + + +class OutputAlignedProjection(NoQuantization): + def __init__(self, alignment=8): + super().__init__() + if type(alignment) is not int or alignment < 1 or alignment & (alignment - 1): + raise ValueError("Projection alignment must be a positive power of two") + self.alignment = alignment + self._padded_weight = None + + def _create_weight(self, out_dims, in_dim, dtype, device_id, num_experts=1): + width = sum(out_dims) if isinstance(out_dims, list) else out_dims + physical_width = (width + self.alignment - 1) // self.alignment * self.alignment + self._padded_weight = None + if dtype not in (torch.float16, torch.bfloat16) or physical_width == width or num_experts != 1: + return super()._create_weight(out_dims, in_dim, dtype, device_id, num_experts) + self._padded_weight = torch.zeros((physical_width, in_dim), dtype=dtype, device=device_id) + pack = WeightPack(weight=self._padded_weight[:width]) + parts = self._split_weight_pack(pack, out_dims if isinstance(out_dims, list) else [out_dims], -2) + return pack, parts + + def apply(self, input_tensor, weight_pack, out=None, workspace=None, use_custom_tensor_mananger=True, bias=None): + if self._padded_weight is None or bias is not None or input_tensor.shape[0] == 0: + return super().apply( + input_tensor, + weight_pack, + out, + workspace, + use_custom_tensor_mananger=use_custom_tensor_mananger, + bias=bias, + ) + from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager + + def allocate(shape): + if use_custom_tensor_mananger: + return g_cache_manager.alloc_tensor(shape, input_tensor.dtype, device=input_tensor.device) + return torch.empty(shape, dtype=input_tensor.dtype, device=input_tensor.device) + + rows, width = input_tensor.shape[0], weight_pack.weight.shape[0] + physical_output = allocate((rows, self._padded_weight.shape[0])) + torch.mm(input_tensor, self._padded_weight.T, out=physical_output) + if out is None: + out = allocate((rows, width)) + elif out.shape != (rows, width) or out.dtype != input_tensor.dtype or out.device != input_tensor.device: + raise ValueError("Projection output must match logical shape, dtype and device") + out.copy_(physical_output[:, :width]) + return out diff --git a/lightllm/common/quantization/scalar_projection.py b/lightllm/common/quantization/scalar_projection.py new file mode 100644 index 0000000000..cef317e4a2 --- /dev/null +++ b/lightllm/common/quantization/scalar_projection.py @@ -0,0 +1,21 @@ +import torch + +from lightllm.common.basemodel.triton_kernel.scalar_linear import scalar_linear +from lightllm.common.quantization.no_quant import NoQuantization + + +class StableScalarProjection(NoQuantization): + def apply(self, input_tensor, weight_pack, out=None, workspace=None, use_custom_tensor_mananger=True, bias=None): + if not input_tensor.is_cuda or weight_pack.weight.shape[0] != 1 or bias is not None: + return super().apply( + input_tensor, + weight_pack, + out=out, + workspace=workspace, + use_custom_tensor_mananger=use_custom_tensor_mananger, + bias=bias, + ) + from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager + + allocate = g_cache_manager.empty if use_custom_tensor_mananger else torch.empty + return scalar_linear(input_tensor, weight_pack.weight, out=out, alloc_func=allocate) diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=2560,N=160,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=2560,N=160,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 0000000000..a4e6e1885e --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=2560,N=160,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,92 @@ +{ + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 3 + }, + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 5 + }, + "3": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 5 + }, + "513": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + }, + "514": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 5 + } +} \ No newline at end of file diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=80,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=80,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 0000000000..444e3b96fe --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=80,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,92 @@ +{ + "88": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 4 + }, + "264": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 4 + }, + "5632": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 4 + }, + "11": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "22": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "33": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 4 + }, + "44": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "513": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "514": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + }, + "5633": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + } +} \ No newline at end of file diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=160,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=160,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H200.json new file mode 100644 index 0000000000..967fd15e61 --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=160,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H200.json @@ -0,0 +1,47 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "1408": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "1409": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "22528": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "22529": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + } +} diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=2560,N=320,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=2560,N=320,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H200.json new file mode 100644 index 0000000000..ace348c88f --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=2560,N=320,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H200.json @@ -0,0 +1,65 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 3 + }, + "33": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + }, + "127": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "513": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "514": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + } +} diff --git a/lightllm/distributed/symm_mem_all_reduce.py b/lightllm/distributed/symm_mem_all_reduce.py index 2256a5093f..7a93318313 100644 --- a/lightllm/distributed/symm_mem_all_reduce.py +++ b/lightllm/distributed/symm_mem_all_reduce.py @@ -90,3 +90,22 @@ def all_reduce(self, inp: torch.Tensor) -> None: else: torch.ops.symm_mem.two_shot_all_reduce_(self.buffer[:n], "sum", self.group.group_name) inp.view(-1).copy_(self.buffer[:n]) + + def ordered_all_gather(self, inp: torch.Tensor): + if not self.should_use(inp) or inp.device != self.device: + return None + count = inp.numel() * self.world_size + if not count or count * inp.element_size() > self.max_size: + return None + if not hasattr(torch.ops.symm_mem, "multimem_all_gather_out"): + return None + if not hasattr(self, "ordered_gather_buffer"): + if torch.cuda.is_current_stream_capturing(): + return None + self.ordered_gather_buffer = torch_symm_mem.empty( + self.max_size // self.dtype.itemsize, device=self.device, dtype=self.dtype + ) + torch_symm_mem.rendezvous(self.ordered_gather_buffer, self.group.group_name) + gathered = self.ordered_gather_buffer[:count] + torch.ops.symm_mem.multimem_all_gather_out(inp.view(-1), self.group.group_name, gathered) + return gathered.view(self.world_size, *inp.shape) diff --git a/lightllm/models/builtin.py b/lightllm/models/builtin.py index d8e091d9aa..cb3bff203c 100644 --- a/lightllm/models/builtin.py +++ b/lightllm/models/builtin.py @@ -209,3 +209,11 @@ def _tarsier_text_model_is(model_type): for config in configs } MODEL_CLASS_PATHS.update({path.rsplit(":", 1)[1]: path for path in DraftModelRegistry._registry.values()}) + +ModelRegistry.register("qwen4_exp", "lightllm.models.qwen4_exp.model:Qwen4ExpTpPartModel", is_multimodal=True) +ModelRegistry.register("qwen4_exp_text", "lightllm.models.qwen4_exp.model:Qwen4ExpTpPartModel") +DraftModelRegistry.register( + ("qwen4_exp", "qwen4_exp_text"), "eagle_with_att", "lightllm.models.qwen4_exp.mtp.model:Qwen4ExpMTPModel" +) +MODEL_CLASS_PATHS["Qwen4ExpTpPartModel"] = "lightllm.models.qwen4_exp.model:Qwen4ExpTpPartModel" +MODEL_CLASS_PATHS["Qwen4ExpMTPModel"] = "lightllm.models.qwen4_exp.mtp.model:Qwen4ExpMTPModel" diff --git a/lightllm/models/glm5_next/config.py b/lightllm/models/glm5_next/config.py index 0e52e736a9..31d79697f6 100644 --- a/lightllm/models/glm5_next/config.py +++ b/lightllm/models/glm5_next/config.py @@ -14,15 +14,17 @@ @dataclasses.dataclass class Glm5NextResourceConfig(Qwen3NextResourceConfig): - """Replicated MLA/index KV plus TP-sharded KDA checkpoints.""" + """描述复制式 MLA/index KV 与 TP 分片 KDA checkpoint 的资源配置。""" # Indexer 池化步长，单位为原始 token；同时决定 K-pool tail 的基准 ring 长度和恢复边界。 + gate_lower_bound: float = -5.0 index_kpool: int = 4 # 单个 indexer key 的特征维度；原始 key/gate tail 的 channel 宽度为该值的两倍。 index_head_dim: int = 128 @classmethod def from_model_config(cls, config: dict[str, object], args: StartArgs | EasyDict) -> "Glm5NextResourceConfig": + """从模型配置和启动参数构造 GLM5.3 Flash 资源配置。""" config = config.get("text_config", config) linear = config["linear_attn_config"] tp = args.tp // args.dp @@ -45,6 +47,7 @@ def from_model_config(cls, config: dict[str, object], args: StartArgs | EasyDict ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), full_attention_interval=4, all_layer_num=layers, + gate_lower_bound=linear.get("gate_lower_bound", -5.0), index_kpool=config["index_kpool"], index_head_dim=config["index_head_dim"], ) diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index e15efe0143..34d28152e8 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -79,6 +79,9 @@ def create_bounded_att_states( request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, + ssm_state_mode=args.ssm_state_mode, + ssm_replay_config=args.ssm_replay_config, + gate_lower_bound=cfg.gate_lower_bound, ) ) full_layers = tuple(model_logical_layer_ids[index] for index in cfg.get_full_att_layer_indices()) @@ -114,37 +117,9 @@ def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, ob config["autotune_layer_num"] = 4 return config - def _init_config(self): - with open(os.path.join(self.weight_dir_, "config.json")) as f: - self.config = self.normalize_model_config(json.load(f)) - - def _verify_params(self): - super()._verify_params() - assert self.config["qk_rope_head_dim"] == 0, "GLM-5.3 Flash uses NoPE attention" - assert not self.args.enable_tpsp_mix_mode, "GLM-5.3 Flash does not support TP/SP mixed mode" - if self.args.llm_kv_type not in ("None", "fp8kv_dsa"): - raise ValueError("GLM-5.3 Flash supports only unquantized KV or dynamic fp8kv_dsa") - if self.args.llm_kv_type == "fp8kv_dsa": - # group64/E8M0 编码与稀疏 attention 算子要求 BF16 输入和固定的 512 维 latent。 - if self.data_type != torch.bfloat16 or self.config["kv_lora_rank"] != 512: - raise ValueError("GLM-5.3 Flash FP8 KV requires BF16 computation and kv_lora_rank=512") - def autotune_layers(self): return 4 - def _init_some_value(self): - self.layers_num = self.config["n_layer"] - self.vocab_size = self.config["vocab_size"] - # MLA stores one replicated latent KV vector per token on every TP rank. - self.tp_k_head_num_ = 1 - self.tp_v_head_num_ = 0 - self.qk_nope_head_dim = self.config["qk_nope_head_dim"] - self.qk_rope_head_dim = self.config["qk_rope_head_dim"] - self.q_lora_rank = self.config["q_lora_rank"] - 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 - @classmethod def create_full_att_states( cls, @@ -155,7 +130,10 @@ 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) @@ -178,18 +156,38 @@ 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 _init_config(self): + with open(os.path.join(self.weight_dir_, "config.json")) as f: + self.config = self.normalize_model_config(json.load(f)) - # Reserve fixed scratch space before the shared KV capacity is profiled. - Fp8Glm5NextSparseAttBackend.allocate_prefill_workspaces(self) + def _verify_params(self): + super()._verify_params() + assert self.config["qk_rope_head_dim"] == 0, "GLM-5.3 Flash uses NoPE attention" + assert not self.args.enable_tpsp_mix_mode, "GLM-5.3 Flash does not support TP/SP mixed mode" + if self.args.llm_kv_type not in ("None", "fp8kv_dsa"): + raise ValueError("GLM-5.3 Flash supports only unquantized KV or dynamic fp8kv_dsa") + if self.args.llm_kv_type == "fp8kv_dsa": + # group64/E8M0 编码与稀疏 attention 算子要求 BF16 输入和固定的 512 维 latent。 + if self.data_type != torch.bfloat16 or self.config["kv_lora_rank"] != 512: + raise ValueError("GLM-5.3 Flash FP8 KV requires BF16 computation and kv_lora_rank=512") + + def _init_some_value(self): + self.layers_num = self.config["n_layer"] + self.vocab_size = self.config["vocab_size"] + # MLA stores one replicated latent KV vector per token on every TP rank. + self.tp_k_head_num_ = 1 + self.tp_v_head_num_ = 0 + self.qk_nope_head_dim = self.config["qk_nope_head_dim"] + self.qk_rope_head_dim = self.config["qk_rope_head_dim"] + self.q_lora_rank = self.config["q_lora_rank"] + 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 + from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import ( + Fp8Glm5NextSparseAttBackend, + ) backend_class = ( Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend @@ -198,7 +196,15 @@ def _init_att_backend(self): self.decode_att_backend = self.prefill_att_backend def _init_att_backend1(self): - self.prefill_att_backend1 = KDALinearAttBackend(model=self) + if getattr(self.args, "ssm_state_mode", "native") == "replay": + from lightllm.common.basemodel.attention.linear.replay_kda import ( + ReplayKDALinearAttBackend, + ) + + backend_cls = ReplayKDALinearAttBackend + else: + backend_cls = KDALinearAttBackend + self.prefill_att_backend1 = backend_cls(model=self) self.decode_att_backend1 = self.prefill_att_backend1 def _init_custom(self): diff --git a/lightllm/models/qwen2_vl/triton_kernel/rotary_pos_emb.py b/lightllm/models/qwen2_vl/triton_kernel/rotary_pos_emb.py index 8063856599..fa61e05c69 100644 --- a/lightllm/models/qwen2_vl/triton_kernel/rotary_pos_emb.py +++ b/lightllm/models/qwen2_vl/triton_kernel/rotary_pos_emb.py @@ -1,4 +1,3 @@ -import math import torch import triton import triton.language as tl @@ -13,6 +12,9 @@ def rotary_kernel( stride_l, stride_h, stride_d, + out_stride_l, + out_stride_h, + out_stride_d, stride_cos_l, stride_cos_d, stride_sin_l, @@ -53,7 +55,7 @@ def rotary_kernel( y = x * cos + rotated * sin - out_ptr_ = out_ptr + base + d * stride_d + out_ptr_ = out_ptr + pid_l * out_stride_l + pid_h * out_stride_h + d * out_stride_d tl.store(out_ptr_, y, mask=mask) @@ -89,6 +91,9 @@ def apply_rotary_pos_emb_triton( stride_l=x.stride(0), stride_h=x.stride(1), stride_d=x.stride(2), + out_stride_l=y.stride(0), + out_stride_h=y.stride(1), + out_stride_d=y.stride(2), stride_cos_l=cos.stride(0), stride_cos_d=cos.stride(1), stride_sin_l=sin.stride(0), diff --git a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py index 1bc4ef914b..df30288562 100644 --- a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py @@ -266,9 +266,7 @@ def context_attention_forward( gdn_out = self._linear_post(core_attn_out, z, layer_weight) - if self.tp_world_size_ > 1: - all_reduce(gdn_out, op=dist.ReduceOp.SUM, group=infer_state.dist_group, async_op=False) - return gdn_out + return self._reduce_linear_output(gdn_out, infer_state) def token_attention_forward( self, @@ -296,9 +294,12 @@ def token_attention_forward( ) gdn_out = self._linear_post(core_attn_out, z, layer_weight) + return self._reduce_linear_output(gdn_out, infer_state) + + def _reduce_linear_output(self, output, infer_state): if self.tp_world_size_ > 1: - all_reduce(gdn_out, op=dist.ReduceOp.SUM, group=infer_state.dist_group, async_op=False) - return gdn_out + all_reduce(output, op=dist.ReduceOp.SUM, group=infer_state.dist_group, async_op=False) + return output def _linear_in_proj( self, @@ -316,7 +317,6 @@ def _linear_post( ) -> torch.Tensor: num_tokens = z.shape[0] core_attn_out = core_attn_out.view(-1, core_attn_out.shape[-1]) - z = z.contiguous().view(-1, z.shape[-1]) norm_out = layer_weight.linear_norm(core_attn_out, z, self.eps_) core_attn_out = norm_out.view(num_tokens, -1) output = layer_weight.linear_out_proj.mm(core_attn_out) diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index f1a3a65748..9da87c6acb 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -81,26 +81,14 @@ def create_bounded_att_states( request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, + ssm_state_mode=args.ssm_state_mode, + ssm_replay_config=args.ssm_replay_config, ) return [state] - def _init_triton(self): - def _triton_allocator(size: int, alignment: int, stream: Optional[int]) -> torch.Tensor: - return torch.empty(size, device="cuda", dtype=torch.int8) - - # Set Triton allocator for TMA descriptors - # This is required for kernels in common/.../linear_att/fla/ops/solve_tril.py - triton.set_allocator(_triton_allocator) - logger.info("Triton allocator set for Qwen3Next model") - return - def autotune_layers(self): return self.config["full_attention_interval"] - def _init_config(self): - super()._init_config() - self.num_kv_heads = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) - @classmethod def create_full_att_states( cls, @@ -135,6 +123,20 @@ def create_full_att_states( ) return [state] + def _init_triton(self): + def _triton_allocator(size: int, alignment: int, stream: Optional[int]) -> torch.Tensor: + return torch.empty(size, device="cuda", dtype=torch.int8) + + # Set Triton allocator for TMA descriptors + # This is required for kernels in common/.../linear_att/fla/ops/solve_tril.py + triton.set_allocator(_triton_allocator) + logger.info("Triton allocator set for Qwen3Next model") + return + + 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): diff --git a/lightllm/models/qwen4_exp/__init__.py b/lightllm/models/qwen4_exp/__init__.py new file mode 100644 index 0000000000..2ae28399f5 --- /dev/null +++ b/lightllm/models/qwen4_exp/__init__.py @@ -0,0 +1 @@ +pass diff --git a/lightllm/models/qwen4_exp/attention_state.py b/lightllm/models/qwen4_exp/attention_state.py new file mode 100644 index 0000000000..d07fa8de89 --- /dev/null +++ b/lightllm/models/qwen4_exp/attention_state.py @@ -0,0 +1,100 @@ +from dataclasses import replace + +import torch + +from lightllm.common.attention_state.full_att.gqa.gqa_state import GQAFullAttentionState, GQAAttentionConfig +from .state import QSAExecutionState + + +class Qwen4ExpAttentionState(GQAFullAttentionState): + def __init__(self, layout_config, indexer_dim, rotary_dim, page_size=1, device="meta"): + super().__init__(layout_config, page_size, device) + common = dict( + global_head_num=1, + dtype=layout_config.dtype, + tp_rank=layout_config.tp_rank, + tp_world_size=layout_config.tp_world_size, + replicated=True, + ) + self.indexer = GQAFullAttentionState( + GQAAttentionConfig(logical_layer_groups=layout_config.logical_layer_groups, head_dim=indexer_dim, **common), + page_size, + device, + ) + layers = tuple(layer for group in layout_config.logical_layer_groups for layer in group) + self.positions = GQAFullAttentionState( + GQAAttentionConfig(logical_layer_groups=(layers,), head_dim=rotary_dim, **common), page_size, device + ) + + def get_state_type(self): + return "qwen4_" + super().get_state_type() + + def allocate_storage(self, token_capacity): + super().allocate_storage(token_capacity) + for state in (self.indexer, self.positions): + state.device = self.device + state.allocate_storage(token_capacity) + + def get_token_bytes(self): + return super().get_token_bytes() + self.indexer.get_token_bytes() + self.positions.get_token_bytes() + + def qsa_view(self, layer): + raw, pooled = self.indexer.get_layer_view(layer) + cos, sin = self.positions.get_layer_view(layer) + return QSAExecutionState(raw[:, 0], pooled[:, 0], cos[:, 0], sin[:, 0]) + + def store_positions(self, indices, cos, sin): + self.positions.k_storage[0, :, 0].index_copy_(0, indices, cos) + self.positions.v_storage[0, :, 0].index_copy_(0, indices, sin) + + def _get_transfer_regions(self, page_tokens): + regions = list(super()._get_transfer_regions(page_tokens)) + for prefix, state in (("indexer", self.indexer), ("positions", self.positions)): + offset = max(region.byte_end for region in regions) + regions.extend( + replace( + region, + field=prefix + "." + region.field, + byte_start=offset + region.byte_start, + byte_end=offset + region.byte_end, + ) + for region in state._get_transfer_regions(page_tokens) + ) + return tuple(regions) + + def _layouts(self, layout): + for prefix, state in (("", super()), ("indexer.", self.indexer), ("positions.", self.positions)): + regions = tuple( + replace(region, field=region.field[len(prefix) :]) + for region in layout.regions + if (region.field.startswith(prefix) if prefix else "." not in region.field) + ) + yield state, replace(layout, regions=regions) + + def write_to_cpu(self, token_indices, cpu_page_buffers, cpu_page_indices, layout, cpu_page_ready_flags=None): + for state, part in self._layouts(layout): + state.write_to_cpu(token_indices, cpu_page_buffers, cpu_page_indices, part, cpu_page_ready_flags) + + def read_from_cpu(self, cpu_page_buffers, cpu_page_indices, token_indices, layout): + for state, part in self._layouts(layout): + state.read_from_cpu(cpu_page_buffers, cpu_page_indices, token_indices, part) + + def write_to_pd_page(self, token_indices, pd_page_buffers, pd_page_index, layout): + for state, part in self._layouts(layout): + state.write_to_pd_page(token_indices, pd_page_buffers, pd_page_index, part) + + def read_from_pd_page(self, pd_page_buffers, token_indices, pd_page_index, layout): + for state, part in self._layouts(layout): + state.read_from_pd_page(pd_page_buffers, token_indices, pd_page_index, part) + + def copy_full_att_tokens(self, src_indices, dst_indices): + super().copy_full_att_tokens(src_indices, dst_indices) + for state in (self.indexer, self.positions): + state.copy_full_att_tokens(src_indices, dst_indices) + + def copy_tokens_from_dp(self, src_token_indices, dst_token_indices, remote_dp_full_att_state): + super().copy_tokens_from_dp(src_token_indices, dst_token_indices, remote_dp_full_att_state) + for name in ("indexer", "positions"): + getattr(self, name).copy_tokens_from_dp( + src_token_indices, dst_token_indices, getattr(remote_dp_full_att_state, name) + ) diff --git a/lightllm/models/qwen4_exp/batch.py b/lightllm/models/qwen4_exp/batch.py new file mode 100644 index 0000000000..299e98e049 --- /dev/null +++ b/lightllm/models/qwen4_exp/batch.py @@ -0,0 +1,119 @@ +from dataclasses import dataclass + +import torch +import triton +import triton.language as tl + +from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_same_req_mark + + +@triton.jit +def _batch_metadata( + Requests, + Lengths, + Starts, + Active, + KeyLengths, + FirstBlocks, + QueryRows, + Visible, + TokenMask, + HOLD: tl.constexpr, + RATIO: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + request = tl.load(Requests + row) + start = tl.load(Starts + row) + end = tl.load(Starts + row + 1) + count = end - start + length = tl.load(Lengths + row) + active = (request != HOLD) & (count > 0) + tl.store(Active + row, tl.where(active, request, -1)) + tl.store(KeyLengths + row, tl.where(active, length, 0)) + tl.store(FirstBlocks + row, tl.where(active, (length - count) // RATIO, 0)) + for first in range(0, tl.cdiv(count, BLOCK)): + offset = first * BLOCK + tl.arange(0, BLOCK) + mask = offset < count + tl.store(QueryRows + start + offset, request, mask) + tl.store(Visible + start + offset, tl.where(active, length - count + offset + 1, 0), mask) + tl.store(TokenMask + start + offset, active, mask) + + +@dataclass(frozen=True) +class PackedBatchMetadata: + starts: torch.Tensor + active_rows: torch.Tensor + key_lengths: torch.Tensor + first_blocks: torch.Tensor + query_rows: torch.Tensor + visible_lengths: torch.Tensor + token_mask: torch.Tensor + + @classmethod + @torch.no_grad() + def prepare( + cls, + requests, + lengths, + query_lengths, + token_count: int, + hold_request: int, + compress_ratio: int, + segment_capacity=None, + ): + if token_count < 0 or compress_ratio < 1: + raise ValueError("Token count must be nonnegative and compression ratio positive") + for tensor in (requests, lengths, query_lengths): + if ( + not tensor.is_cuda + or tensor.device != requests.device + or tensor.dtype != torch.int32 + or tensor.ndim != 1 + or not tensor.is_contiguous() + or tensor.shape != requests.shape + ): + raise ValueError("Batch metadata requires matching contiguous int32 CUDA vectors") + batch = requests.numel() + capacity = batch if segment_capacity is None else segment_capacity + if capacity < batch: + raise ValueError("Metadata segment capacity cannot be smaller than the input batch") + if not batch and token_count: + raise ValueError("An empty batch cannot contain input tokens") + + starts = torch.full((capacity + 1,), token_count, dtype=torch.int32, device=requests.device) + starts[:1].zero_() + if batch: + marks = gen_same_req_mark(requests, hold_req_idx=hold_request) + counts = torch.where(marks > 0, query_lengths + marks - 1, 0) + torch.cumsum(counts, dim=0, dtype=torch.int32, out=starts[1 : batch + 1]) + active = torch.full((capacity,), -1, dtype=torch.int32, device=requests.device) + key_lengths = torch.zeros_like(active) + first_blocks = torch.zeros_like(active) + query_rows = torch.empty(token_count, dtype=torch.int32, device=requests.device) + visible = torch.empty_like(query_rows) + token_mask = torch.empty(token_count, dtype=torch.bool, device=requests.device) + if batch: + _batch_metadata[(batch,)]( + requests, + lengths, + starts, + active, + key_lengths, + first_blocks, + query_rows, + visible, + token_mask, + hold_request, + compress_ratio, + 256, + ) + return cls(starts, active, key_lengths, first_blocks, query_rows, visible, token_mask) + + def copy_from(self, other: "PackedBatchMetadata"): + for name in self.__dataclass_fields__: + target, source = getattr(self, name), getattr(other, name) + if target.shape != source.shape or target.device != source.device or target.dtype != source.dtype: + raise ValueError(f"Incompatible graph batch metadata: {name}") + for name in self.__dataclass_fields__: + getattr(self, name).copy_(getattr(other, name), non_blocking=True) diff --git a/lightllm/models/qwen4_exp/checkpoint.py b/lightllm/models/qwen4_exp/checkpoint.py new file mode 100644 index 0000000000..e5a42a716b --- /dev/null +++ b/lightllm/models/qwen4_exp/checkpoint.py @@ -0,0 +1,38 @@ +import math +from functools import lru_cache + + +def normalize_text_weight_names(weights: dict): + prefix = "model.language_model." + for name in tuple(weights): + if name.startswith(prefix): + target = "model." + name[len(prefix) :] + if target in weights and weights[target] is not weights[name]: + raise ValueError(f"Conflicting text checkpoint names: {name} and {target}") + weights[target] = weights.pop(name) + + +def ngram_eos_token_id(config: dict) -> int: + eos = config["eos_token_id"] + if isinstance(eos, list): + if not eos: + raise ValueError("PLE requires an EOS token ID") + eos = eos[0] + return eos + + +@lru_cache(maxsize=None) +def ngram_table_shape(base: int, heads: int, ple_index: int, embedding_dim: int, padding: int): + if base < 1 or heads < 1 or ple_index < 0 or embedding_dim < 1 or embedding_dim % heads or padding < 1: + raise ValueError("Invalid n-gram table dimensions") + sizes = [] + candidate = base - 1 + for index in range((ple_index + 1) * heads): + while True: + candidate += 1 + if candidate >= 2 and all(candidate % divisor for divisor in range(2, math.isqrt(candidate) + 1)): + break + if index >= ple_index * heads: + sizes.append(candidate) + total = sum(sizes) + return ((total + padding - 1) // padding) * padding, embedding_dim // heads diff --git a/lightllm/models/qwen4_exp/configuration.py b/lightllm/models/qwen4_exp/configuration.py new file mode 100644 index 0000000000..a47e250b84 --- /dev/null +++ b/lightllm/models/qwen4_exp/configuration.py @@ -0,0 +1,24 @@ +from lightllm.common.build_utils import repair_config + + +def normalize_model_config(checkpoint, finetune_config=None): + config = dict(checkpoint.get("text_config", checkpoint)) + for aliases in ( + ["num_attention_heads", "n_head"], + ["hidden_size", "n_embd", "n_embed"], + ["num_hidden_layers", "n_layer"], + ): + repair_config(config, same_names=aliases) + rope = config.get("rope_parameters") + if isinstance(rope, dict): + for name in ("rope_theta", "partial_rotary_factor"): + if name in rope: + config.setdefault(name, rope[name]) + config.setdefault("rope_scaling", dict(rope)) + for name in ("image_token_id", "video_token_id", "vision_start_token_id", "vision_end_token_id"): + if name in checkpoint: + config[name] = checkpoint[name] + config.setdefault("norm_topk_prob", True) + if finetune_config is not None: + config["vocab_size"] = finetune_config.vocab_size + return config, checkpoint.get("vision_config") diff --git a/lightllm/models/qwen4_exp/image_processing.py b/lightllm/models/qwen4_exp/image_processing.py new file mode 100644 index 0000000000..902e0d9d36 --- /dev/null +++ b/lightllm/models/qwen4_exp/image_processing.py @@ -0,0 +1,9 @@ +from lightllm.models.qwen2_vl.vision_process import Qwen2VLImageProcessor + + +class Qwen4ExpImageProcessor(Qwen2VLImageProcessor): + @classmethod + def from_pretrained(cls, model_path, **kwargs): + config, overrides = cls.get_image_processor_dict(model_path, **kwargs) + config.update(overrides) + return cls(**config) diff --git a/lightllm/models/qwen4_exp/infer_struct.py b/lightllm/models/qwen4_exp/infer_struct.py new file mode 100644 index 0000000000..06062d7959 --- /dev/null +++ b/lightllm/models/qwen4_exp/infer_struct.py @@ -0,0 +1,76 @@ +from lightllm.models.qwen3_5.infer_struct import Qwen35InferStateInfo +from .batch import PackedBatchMetadata +from .position import combine_mrope +from .state import QSABatch + + +class Qwen4ExpInferStateInfo(Qwen35InferStateInfo): + def init_some_extra_state(self, model): + super().init_some_extra_state(model) + config = model.config + ratio = config["indexer_compress_ratio"] + self.qwen4_exp_metadata = PackedBatchMetadata.prepare( + self.b_req_idx, + self.b_seq_len, + self.b_q_seq_len, + self.input_ids.numel(), + self.req_manager.HOLD_REQUEST_ID, + ratio, + segment_capacity=model.max_req_num + 1 if self.is_prefill else None, + ) + cos, sin = combine_mrope(self.position_cos, self.position_sin, config["rope_scaling"]["mrope_section"]) + self.qwen4_exp_batch = QSABatch.from_metadata( + self.qwen4_exp_metadata, self.req_manager.req_to_token_indexs, self.mem_index, cos, sin, ratio + ) + from .attention_state import Qwen4ExpAttentionState + from .lexical_state import Qwen4ExpLexicalState + + states = self.att_state_manager.get_full_attention_states(self.att_state_role) + self.qwen4_exp_qsa = {} + for state in states: + if isinstance(state, Qwen4ExpAttentionState): + state.store_positions(self.qwen4_exp_batch.mem_indices, cos, sin) + for group in state.config.logical_layer_groups: + for layer in group: + self.qwen4_exp_qsa[layer] = state.qsa_view(layer) + self.qwen4_exp_lexical_state = next( + ( + state + for state in self.att_state_manager.get_bounded_attention_states(self.att_state_role) + if isinstance(state, Qwen4ExpLexicalState) + ), + None, + ) + self.qwen4_exp_ple = ( + {} + if self.qwen4_exp_lexical_state is None + else self.qwen4_exp_lexical_state.prepare_lexical( + self.qwen4_exp_metadata, reset_empty_prefix=self.is_prefill + ) + ) + self.qwen4_exp_auto_commit = self.is_prefill or model.args.mtp_step == 0 + + def commit_lexical(self, accepted_lengths): + self.qwen4_exp_lexical_state.commit_lexical(self.qwen4_exp_metadata, self.qwen4_exp_ple, accepted_lengths) + + def _copy_qwen_state(self, new_state): + if self.qwen4_exp_auto_commit != new_state.qwen4_exp_auto_commit: + raise ValueError("Qwen graph replay cannot change lexical commit mode") + self.qwen4_exp_metadata.copy_from(new_state.qwen4_exp_metadata) + old, new = self.qwen4_exp_batch, new_state.qwen4_exp_batch + if old.max_new_blocks != new.max_new_blocks: + raise ValueError("Qwen graph replay requires the same token capacity") + for name in ("mem_indices", "cos", "sin"): + getattr(old, name).copy_(getattr(new, name), non_blocking=True) + for layer, batch in self.qwen4_exp_ple.items(): + source = new_state.qwen4_exp_ple[layer] + batch.ngram_history.copy_(source.ngram_history, non_blocking=True) + batch.conv_history.copy_(source.conv_history, non_blocking=True) + + def copy_for_cuda_graph(self, new_infer_state): + super().copy_for_cuda_graph(new_infer_state) + self._copy_qwen_state(new_infer_state) + + def copy_for_prefill_cuda_graph(self, new_infer_state): + super().copy_for_prefill_cuda_graph(new_infer_state) + self._copy_qwen_state(new_infer_state) diff --git a/lightllm/models/qwen4_exp/layer_infer/__init__.py b/lightllm/models/qwen4_exp/layer_infer/__init__.py new file mode 100644 index 0000000000..2ae28399f5 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_infer/__init__.py @@ -0,0 +1 @@ +pass diff --git a/lightllm/models/qwen4_exp/layer_infer/post_layer_infer.py b/lightllm/models/qwen4_exp/layer_infer/post_layer_infer.py new file mode 100644 index 0000000000..a685363b1d --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_infer/post_layer_infer.py @@ -0,0 +1,6 @@ +from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer + + +class Qwen4ExpPostLayerInfer(LlamaPostLayerInfer): + def _norm(self, input, infer_state, layer_weight): + return layer_weight.hyper_connection_mixer.component.read(input) diff --git a/lightllm/models/qwen4_exp/layer_infer/pre_layer_infer.py b/lightllm/models/qwen4_exp/layer_infer/pre_layer_infer.py new file mode 100644 index 0000000000..129ab20853 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_infer/pre_layer_infer.py @@ -0,0 +1,23 @@ +from lightllm.models.qwen3_vl.layer_infer.pre_layer_infer import Qwen3VLMultimodalPreLayerInfer +from ..multimodal import canonical_image_ids + + +class Qwen4ExpPreLayerInfer(Qwen3VLMultimodalPreLayerInfer): + def __init__(self, network_config): + super().__init__(network_config) + self.branches = network_config["hc_count"] + self.image_token_id = network_config.get("image_token_id") + + def context_forward(self, input_ids, infer_state, layer_weight): + if not any(params.get("images") or params.get("audios") for params in (infer_state.multimodal_params or [])): + return self.token_forward(input_ids, infer_state, layer_weight) + embeddings = super().context_forward(input_ids, infer_state, layer_weight) + infer_state.qwen4_exp_lexical_input_ids = canonical_image_ids( + input_ids, infer_state.img_start_token_ids, infer_state.img_token_lens, self.image_token_id + ) + return embeddings.repeat(1, self.branches) + + def token_forward(self, input_ids, infer_state, layer_weight): + infer_state.qwen4_exp_lexical_input_ids = input_ids + embeddings = super().token_forward(input_ids, infer_state, layer_weight) + return embeddings.repeat(1, self.branches) diff --git a/lightllm/models/qwen4_exp/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen4_exp/layer_infer/transformer_layer_infer.py new file mode 100644 index 0000000000..7ccb2e98d3 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_infer/transformer_layer_infer.py @@ -0,0 +1,93 @@ +from lightllm.common.basemodel.tp_reduce import ordered_tp_sum +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.models.qwen3_5.layer_infer.transformer_layer_infer import Qwen35TransformerLayerInfer + + +class Qwen4ExpTransformerLayerInfer(Qwen35TransformerLayerInfer): + def __init__(self, layer_num, network_config): + super().__init__(layer_num, network_config) + layer_types = network_config.get("layer_types") + if layer_types is not None: + self.is_linear_attention_layer = layer_types[layer_num] == "linear_attention" + + def _reduce_linear_output(self, output, infer_state): + return self._tpsp_reduce(output, infer_state) + + def _tpsp_reduce(self, input, infer_state): + if self.tp_world_size_ == 1: + return input + scatter_rank = self.tp_rank_ if get_env_start_args().enable_tpsp_mix_mode else None + return ordered_tp_sum(input, infer_state.dist_group, self.tp_world_size_, scatter_rank=scatter_rank) + + def _forward(self, residual, infer_state, layer_weight, is_prefill): + if hasattr(layer_weight, "ple"): + ple_batch = infer_state.qwen4_exp_ple[self.layer_num_] + lexical_tokens = getattr(infer_state, "qwen4_exp_lexical_input_ids", infer_state.input_ids) + lexical = ple_batch.apply( + self._tpsp_allgather(residual, infer_state), + lexical_tokens, + layer_weight, + self.tp_world_size_, + infer_state.dist_group, + ) + auto_commit = getattr(infer_state, "qwen4_exp_auto_commit", None) + if auto_commit: + infer_state.qwen4_exp_lexical_state.lexical_caches[self.layer_num_].commit( + infer_state.qwen4_exp_metadata, ple_batch + ) + elif auto_commit is False: + infer_state.qwen4_exp_lexical_state.stage_speculative_lexical( + self.layer_num_, infer_state.qwen4_exp_metadata, ple_batch + ) + residual = residual + self._tpsp_sp_split(lexical, infer_state) + attention_residual = layer_weight.attn_hyper_connection.component + mixed, gates = attention_residual.read(residual) + if is_prefill: + update = self.context_attention_forward(mixed, infer_state, layer_weight) + else: + update = self.token_attention_forward(mixed, infer_state, layer_weight) + residual = attention_residual.write(residual, update, gates) + mlp_residual = layer_weight.mlp_hyper_connection.component + mixed, gates = mlp_residual.read(residual) + update = self._ffn(mixed, infer_state, layer_weight) + return mlp_residual.write(residual, update, gates) + + def context_forward(self, input_embdings, infer_state, layer_weight): + return self._forward(input_embdings, infer_state, layer_weight, True) + + def token_forward(self, input_embdings, infer_state, layer_weight): + return self._forward(input_embdings, infer_state, layer_weight, False) + + def _qsa_attention_forward(self, input_embdings, infer_state, layer_weight): + query, kv = self._get_qkv(input_embdings, infer_state, layer_weight) + self._post_cache_kv(kv, infer_state, layer_weight) + from ..attention_state import Qwen4ExpAttentionState + + state = infer_state.att_state_manager.get_state( + infer_state.att_state_role, "qwen4_gqa", self.layer_num_, Qwen4ExpAttentionState + ) + key, value = state.get_layer_view(self.layer_num_) + hidden = self._tpsp_allgather(input_embdings, infer_state) + output = infer_state.qwen4_exp_qsa[self.layer_num_].attend( + layer_weight.indexer.component, + hidden, + query.view(-1, self.tp_q_head_num_, self.head_dim_), + key, + value, + infer_state.qwen4_exp_batch, + decode=not infer_state.is_prefill, + attention_splits=None if infer_state.is_prefill else 1, + ) + return self._get_o(output, infer_state, layer_weight) + + def context_attention_forward(self, input_embdings, infer_state, layer_weight): + if self.is_linear_attention_layer: + input_embdings = self._tpsp_allgather(input_embdings, infer_state) + return super().context_attention_forward(input_embdings, infer_state, layer_weight) + return self._qsa_attention_forward(input_embdings, infer_state, layer_weight) + + def token_attention_forward(self, input_embdings, infer_state, layer_weight): + if self.is_linear_attention_layer: + input_embdings = self._tpsp_allgather(input_embdings, infer_state) + return super().token_attention_forward(input_embdings, infer_state, layer_weight) + return self._qsa_attention_forward(input_embdings, infer_state, layer_weight) diff --git a/lightllm/models/qwen4_exp/layer_weights/__init__.py b/lightllm/models/qwen4_exp/layer_weights/__init__.py new file mode 100644 index 0000000000..2ae28399f5 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_weights/__init__.py @@ -0,0 +1 @@ +pass diff --git a/lightllm/models/qwen4_exp/layer_weights/components.py b/lightllm/models/qwen4_exp/layer_weights/components.py new file mode 100644 index 0000000000..8c0e9d5717 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_weights/components.py @@ -0,0 +1,129 @@ +from functools import partial +from typing import Callable, Mapping + +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights.base_weight import BaseWeight +from lightllm.common.basemodel.layer_weights.meta_weights.parameter_weight import ParameterWeight +from lightllm.common.gated_residual import GatedResidual +from lightllm.common.ngram_embedding import NGramHasher +from lightllm.common.ngram_embedding.ple import PLELayer +from lightllm.common.qsa.indexer import QSAIndexer +from lightllm.models.qwen4_exp.checkpoint import ngram_eos_token_id + + +class ComponentWeight(BaseWeight): + def __init__(self, specs: Mapping[str, tuple], builder: Callable): + super().__init__() + self._specs = specs + self._builder = builder + self._component = None + self._create_weight() + + def _create_weight(self): + self.parameters = {name: ParameterWeight(name, dtype, shape) for name, (shape, dtype) in self._specs.items()} + + def load_hf_weights(self, weights): + for name, parameter in self.parameters.items(): + if name not in weights: + continue + source = weights[name] + shape, dtype = self._specs[name] + if source.shape != shape: + raise ValueError(f"Component checkpoint shape mismatch for {name}: {tuple(source.shape)} != {shape}") + if dtype == torch.int64 and source.dtype != torch.int64: + raise ValueError(f"Integer component buffer {name} must use int64") + if dtype != torch.int64 and source.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError(f"Component tensor {name} requires unquantized floating point weights") + parameter.load_hf_weights(weights) + + def verify_load(self): + ready = all(parameter.verify_load() for parameter in self.parameters.values()) + if ready and self._component is None: + self._component = self._builder({name: parameter.weight for name, parameter in self.parameters.items()}) + return ready + + @property + def component(self): + if self._component is None and not self.verify_load(): + raise RuntimeError("Component weights are not fully loaded") + return self._component + + +def gated_residual_weight(prefix: str, config: dict, dtype: torch.dtype, use_combine: bool = True): + width = config["hc_count"] * config["hidden_size"] + rank = config["hc_lowrank"] + specs = { + f"{prefix}.hc_norm.weight": ((width,), dtype), + f"{prefix}.input_mix_weight_down.weight": ((rank, width), dtype), + f"{prefix}.input_mix_weight_up.weight": ((width, rank), dtype), + } + if use_combine: + specs[f"{prefix}.block_inject_weight.weight"] = ((config["hc_count"], width), dtype) + return ComponentWeight( + specs, + partial( + GatedResidual.from_checkpoint, + prefix=prefix, + branches=config["hc_count"], + eps=config["rms_norm_eps"], + use_combine=use_combine, + ), + ) + + +def qsa_indexer_weight(prefix: str, config: dict, dtype: torch.dtype): + dim, heads = config["indexer_head_dim"], config["indexer_n_heads"] + if config["indexer_kv_heads"] != 1: + raise ValueError("Qwen4Exp QSA requires one index key head") + specs = { + f"{prefix}.index_qk_proj.weight": (((heads + 1) * dim, config["hidden_size"]), dtype), + f"{prefix}.q_layernorm.weight": ((dim,), dtype), + f"{prefix}.k_layernorm.weight": ((dim,), dtype), + } + return ComponentWeight( + specs, + partial( + QSAIndexer.from_checkpoint, + prefix=prefix, + query_heads=heads, + token_budget=config["indexer_budget"], + use_cuda_topk=True, + compress_ratio=config["indexer_compress_ratio"], + eps=config["rms_norm_eps"], + ), + ) + + +def ple_weight(prefix: str, config: dict, dtype: torch.dtype): + width = config["hc_count"] * config["hidden_size"] + embed = config.get("ple_embed_dim", config["hidden_size"]) + specs = { + f"{prefix}.key_proj.weight": ((width, embed), dtype), + f"{prefix}.value_proj.weight": ((config["hidden_size"], embed), dtype), + f"{prefix}.conv1d.weight": ((width, 1, config["ple_conv_kernel_size"]), dtype), + } + for norm in ("norm_key", "norm_query", "norm_conv"): + specs[f"{prefix}.{norm}.weight"] = ((width,), dtype) + return ComponentWeight( + specs, + partial( + PLELayer.from_checkpoint, + prefix=prefix, + branches=config["hc_count"], + dilation=config["ngram_size"], + eps=config["rms_norm_eps"], + ), + ) + + +def ngram_hasher_weight(prefix: str, config: dict): + heads = (config["ngram_size"] - 1) * config["heads_per_ngram"] + specs = { + f"{prefix}.layer_multipliers": ((config["ngram_size"],), torch.int64), + f"{prefix}.ngram_heads_vocab_sizes": ((heads,), torch.int64), + f"{prefix}.ngram_heads_offsets": ((heads,), torch.int64), + } + return ComponentWeight( + specs, partial(NGramHasher.from_checkpoint, prefix=prefix, eos_token_id=ngram_eos_token_id(config)) + ) diff --git a/lightllm/models/qwen4_exp/layer_weights/embedding_weight.py b/lightllm/models/qwen4_exp/layer_weights/embedding_weight.py new file mode 100644 index 0000000000..ef2f5980dd --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_weights/embedding_weight.py @@ -0,0 +1,14 @@ +from lightllm.common.basemodel.layer_weights.meta_weights.row_sharded_embedding_weight import RowShardedEmbeddingWeight +from lightllm.common.basemodel.layer_weights.meta_weights.pinned_embedding_weight import PinnedRowShardedEmbeddingWeight +from lightllm.utils.envs_utils import get_env_start_args + + +def ngram_embedding_weight(dim, rows, name, dtype, shard_count): + args = get_env_start_args() + if not args.get("offload_ngram_embedding", False): + return RowShardedEmbeddingWeight(dim, rows, name, dtype, shard_count) + if args.get("hardware_platform", "cuda") != "cuda": + raise ValueError("N-gram embedding offload requires CUDA Triton lookup") + from ..pinned_memory import pinned_zeros + + return PinnedRowShardedEmbeddingWeight(dim, rows, name, dtype, shard_count, cpu_allocator=pinned_zeros) diff --git a/lightllm/models/qwen4_exp/layer_weights/lm_head_weight.py b/lightllm/models/qwen4_exp/layer_weights/lm_head_weight.py new file mode 100644 index 0000000000..72dd4d1f50 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_weights/lm_head_weight.py @@ -0,0 +1,25 @@ +from typing import Optional + +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights import LMHeadWeight + + +class Qwen4ExpLMHeadWeight(LMHeadWeight): + def _cuda_forward( + self, input: torch.Tensor, out: Optional[torch.Tensor] = None, alloc_func=torch.empty + ) -> torch.Tensor: + rows = input.shape[1] + if rows <= 1 or rows % 8 == 0 or input.dtype not in (torch.float16, torch.bfloat16): + return super()._cuda_forward(input, out=out, alloc_func=alloc_func) + + tokens = input.transpose(0, 1) + padded = torch.cat((tokens, tokens[:1].expand((-rows) % 8, -1)), dim=0) + scores = super()._cuda_forward(padded.transpose(0, 1), alloc_func=alloc_func) + shape = (self.weight.shape[0], rows) + if out is None: + out = alloc_func(shape, dtype=input.dtype, device=input.device) + elif out.shape != shape or out.dtype != input.dtype or out.device != input.device: + raise ValueError("LM-head output must match the unpadded shape, dtype and device") + out.copy_(scores[:, :rows]) + return out diff --git a/lightllm/models/qwen4_exp/layer_weights/norm_weight.py b/lightllm/models/qwen4_exp/layer_weights/norm_weight.py new file mode 100644 index 0000000000..1824e4948a --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_weights/norm_weight.py @@ -0,0 +1,30 @@ +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights import GatedRMSNormWeight +from lightllm.models.qwen4_exp.triton_kernel.gated_norm import gated_norm + + +class Qwen4ExpGatedNormWeight(GatedRMSNormWeight): + def __init__(self, dim, weight_name, data_type, activation="sigmoid"): + if activation not in ("sigmoid", "silu"): + raise ValueError("Qwen4Exp GDN output gate must be sigmoid or silu") + self.activation = activation + super().__init__(dim, weight_name, data_type) + + def _triton_forward(self, input, gate_value, eps, out=None, alloc_func=torch.empty): + if out is None: + out = alloc_func(input.shape, dtype=input.dtype, device=input.device) + return gated_norm(input, gate_value, self.weight, eps, self.activation, out) + + def _native_forward(self, input, gate_value, eps, out=None, alloc_func=torch.empty): + values = input.float() + normalized = values * torch.rsqrt(values.square().mean(-1, keepdim=True) + eps) + weighted = self.weight * normalized.to(input.dtype) + gate = gate_value.float().sigmoid() + if self.activation == "silu": + gate = gate * gate_value.float() + result = (weighted * gate).to(input.dtype) + if out is not None: + out.copy_(result) + return out + return result diff --git a/lightllm/models/qwen4_exp/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/qwen4_exp/layer_weights/pre_and_post_layer_weight.py new file mode 100644 index 0000000000..dabedaed88 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_weights/pre_and_post_layer_weight.py @@ -0,0 +1,26 @@ +from lightllm.common.basemodel import PreAndPostLayerWeight +from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight +from lightllm.models.qwen4_exp.checkpoint import normalize_text_weight_names +from .components import gated_residual_weight +from .lm_head_weight import Qwen4ExpLMHeadWeight + + +class Qwen4ExpPreAndPostLayerWeight(PreAndPostLayerWeight): + def __init__(self, data_type, network_config): + super().__init__(data_type, network_config) + hidden, vocab = network_config["hidden_size"], network_config["vocab_size"] + self.wte_weight_ = EmbeddingWeight(hidden, vocab, "model.embed_tokens.weight", data_type) + self.lm_head_weight_ = Qwen4ExpLMHeadWeight( + hidden, + vocab, + "lm_head.weight", + data_type, + self.wte_weight_ if network_config.get("tie_word_embeddings", False) else None, + ) + self.hyper_connection_mixer = gated_residual_weight( + "model.hyper_connection_mixer", network_config, data_type, use_combine=False + ) + + def load_hf_weights(self, weights): + normalize_text_weight_names(weights) + super().load_hf_weights(weights) diff --git a/lightllm/models/qwen4_exp/layer_weights/transformer_layer_weight.py b/lightllm/models/qwen4_exp/layer_weights/transformer_layer_weight.py new file mode 100644 index 0000000000..25e2d082b2 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_weights/transformer_layer_weight.py @@ -0,0 +1,85 @@ +import torch + +from lightllm.models.qwen3_5_moe.layer_weights.transformer_layer_weight import Qwen35MOETransformerLayerWeight +from lightllm.common.basemodel.layer_weights.meta_weights import QKGEMMANormWeight +from lightllm.common.quantization.scalar_projection import StableScalarProjection +from lightllm.common.quantization.aligned_projection import OutputAlignedProjection +from lightllm.common.quantization.no_quant import NoQuantization +from lightllm.models.qwen4_exp.checkpoint import normalize_text_weight_names, ngram_table_shape +from .components import gated_residual_weight, qsa_indexer_weight, ple_weight, ngram_hasher_weight +from .norm_weight import Qwen4ExpGatedNormWeight +from .embedding_weight import ngram_embedding_weight + + +class Qwen4ExpTransformerLayerWeight(Qwen35MOETransformerLayerWeight): + def get_quant_method(self, name): + method = super().get_quant_method(name) + if name == "in_proj_weight" and type(method) is NoQuantization: + aligned = OutputAlignedProjection() + aligned.hf_quantization_config = method.hf_quantization_config + return aligned + return method + + def _init_moe(self): + super()._init_moe() + + self.shared_expert_gate.quant_method = StableScalarProjection() + + def _init_weight(self): + config = self.network_config_ + layer_types = config.get("layer_types") + if layer_types is not None: + layer_type = layer_types[self.layer_num_] + if layer_type not in ("linear_attention", "full_attention", "indexed_attention"): + raise ValueError(f"Unsupported Qwen4Exp layer type {layer_type}") + self.is_linear_attention_layer = layer_type == "linear_attention" + super()._init_weight() + prefix = f"model.layers.{self.layer_num_}" + self.attn_hyper_connection = gated_residual_weight(f"{prefix}.attn_hyper_connection", config, self.data_type_) + self.mlp_hyper_connection = gated_residual_weight(f"{prefix}.mlp_hyper_connection", config, self.data_type_) + if not self.is_linear_attention_layer: + self.indexer = qsa_indexer_weight(f"{prefix}.self_attn.indexer", config, self.data_type_) + ple_layers = sorted(set(config.get("ple_layer_ids", []))) + if self.layer_num_ + 1 in ple_layers: + if not self.is_linear_attention_layer: + raise ValueError("Qwen4Exp PLE must be attached to a linear-attention layer") + self.ple = ple_weight(f"{prefix}.ple", config, self.data_type_) + self.ngram_hasher = ngram_hasher_weight(f"{prefix}.ple.ple_embedding", config) + heads = (config["ngram_size"] - 1) * config["heads_per_ngram"] + rows, dim = ngram_table_shape( + config["ngram_vocab_size_base"], + heads, + ple_layers.index(self.layer_num_ + 1), + config.get("ple_embed_dim", config["hidden_size"]), + config["make_ngram_vocab_size_divisible_by"], + ) + self.ngram_embedding = ngram_embedding_weight( + dim, + rows, + f"{prefix}.ple.ple_embedding.ngram_embedding.weight", + self.data_type_, + config["split_ngram_parts"], + ) + + def _init_gdn_weight(self): + super()._init_gdn_weight() + config = self.network_config_ + self.linear_norm = Qwen4ExpGatedNormWeight( + self.linear_v_head_dim, + f"model.layers.{self.layer_num_}.linear_attn.norm.weight", + self.data_type_, + config.get("output_gate_type") or config["hidden_act"], + ) + + def _init_norm(self): + if not self.is_linear_attention_layer: + self.qk_norm_weight_ = QKGEMMANormWeight( + self.head_dim, + self._q_norm_name, + self._k_norm_name, + torch.float32, + ) + + def load_hf_weights(self, weights): + normalize_text_weight_names(weights) + super().load_hf_weights(weights) diff --git a/lightllm/models/qwen4_exp/lexical_cache.py b/lightllm/models/qwen4_exp/lexical_cache.py new file mode 100644 index 0000000000..a5cdedab4b --- /dev/null +++ b/lightllm/models/qwen4_exp/lexical_cache.py @@ -0,0 +1,89 @@ +from dataclasses import dataclass + +import torch +import triton +import triton.language as tl + +from .state import PLEBatch +from .checkpoint import ngram_eos_token_id +from lightllm.utils.dist_utils import get_current_device_id + + +@triton.jit +def _scatter_history(Source, Destinations, Target, WIDTH: tl.constexpr, BLOCK: tl.constexpr): + row = tl.program_id(0) + destination = tl.load(Destinations + row).to(tl.int64) + if destination >= 0: + col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + values = tl.load(Source + row.to(tl.int64) * WIDTH + col, col < WIDTH, other=0) + tl.store(Target + destination * WIDTH + col, values, col < WIDTH) + + +@dataclass(frozen=True) +class LexicalSnapshot: + ngram: torch.Tensor + convolution: torch.Tensor + + +class LexicalRequestCache: + def __init__(self, layer_weight, requests: int): + if requests < 1: + raise ValueError("Lexical cache requires at least one request slot") + self.layer_weight = layer_weight + config = layer_weight.network_config_ + self.eos_token_id = ngram_eos_token_id(config) + + device = torch.device("cuda", get_current_device_id()) + self.ngram = torch.full( + (requests, config["ngram_size"] - 1), self.eos_token_id, dtype=torch.int64, device=device + ) + self.convolution = torch.zeros( + ( + requests, + (config["ple_conv_kernel_size"] - 1) * config["ngram_size"], + config["hc_count"] * config["hidden_size"], + ), + dtype=layer_weight.data_type_, + device=device, + ) + + def reset(self, request: int): + self.ngram[request].fill_(self.eos_token_id) + self.convolution[request].zero_() + + @torch.no_grad() + def prepare(self, metadata, reset_empty_prefix=False) -> PLEBatch: + rows = metadata.active_rows + indices = rows.clamp_min(0).to(torch.int64) + ngram = self.ngram.index_select(0, indices) + convolution = self.convolution.index_select(0, indices) + empty = rows < 0 + if reset_empty_prefix: + counts = metadata.starts[1:] - metadata.starts[:-1] + empty = empty | (metadata.key_lengths == counts) + ngram.masked_fill_(empty[:, None], self.eos_token_id) + convolution.masked_fill_(empty[:, None, None], 0) + return PLEBatch(metadata.starts, ngram, convolution, metadata.token_mask) + + @torch.no_grad() + def commit(self, metadata, batch: PLEBatch, accepted_lengths=None): + ngram, convolution = batch.histories_after(self.layer_weight, accepted_lengths) + for values, storage in ((ngram, self.ngram), (convolution, self.convolution)): + width = storage[0].numel() + if metadata.active_rows.numel() and width: + _scatter_history[(metadata.active_rows.numel(), triton.cdiv(width, 256))]( + values, metadata.active_rows, storage, width, 256 + ) + + def snapshot(self, request: int) -> LexicalSnapshot: + return LexicalSnapshot(self.ngram[request].clone(), self.convolution[request].clone()) + + def restore(self, request: int, snapshot: LexicalSnapshot): + for target, source in ( + (self.ngram[request], snapshot.ngram), + (self.convolution[request], snapshot.convolution), + ): + if target.shape != source.shape or target.dtype != source.dtype: + raise ValueError("Lexical snapshot has incompatible shape or dtype") + self.ngram[request].copy_(snapshot.ngram, non_blocking=True) + self.convolution[request].copy_(snapshot.convolution, non_blocking=True) diff --git a/lightllm/models/qwen4_exp/lexical_state.py b/lightllm/models/qwen4_exp/lexical_state.py new file mode 100644 index 0000000000..cf13172468 --- /dev/null +++ b/lightllm/models/qwen4_exp/lexical_state.py @@ -0,0 +1,257 @@ +from dataclasses import dataclass +from functools import lru_cache +from math import prod + +import torch + +from lightllm.common.attention_state.bounded_att import BoundedAttentionState +from lightllm.common.attention_state.bounded_att.copy_kernel import copy_bounded_att_tensors +from lightllm.common.attention_state.bounded_att.transfer_utils import bounded_att_buffer_copies +from lightllm.common.attention_state.bounded_att.utils import create_page_pool_tensor, register_page_pool_tensor +from lightllm.common.attention_state.types import AxisMapping, TensorField, TransferRegion +from lightllm.common.attention_state.utils import deepcopy_result +from .lexical_cache import LexicalRequestCache +from .speculative_cache import LexicalSpeculativeCache + + +@dataclass(frozen=True) +class LexicalStateConfig: + logical_layer_ids: tuple[int, ...] + ngram_context: int + conv_context: int + width: int + eos_token_id: int + dtype: torch.dtype + tp_rank: int = 0 + + +class LexicalPages: + def __init__(self, capacity, config, device, pin_memory): + count = len(config.logical_layer_ids) + self.ngram = create_page_pool_tensor( + size=(capacity, count, config.ngram_context), + dtype=torch.int64, + device=torch.device(device), + pin_memory=pin_memory, + ) + self.convolution = create_page_pool_tensor( + size=(capacity, count, config.conv_context, config.width), + dtype=config.dtype, + device=torch.device(device), + pin_memory=pin_memory, + ) + + def register_pin_memory(self): + for tensor in (self.ngram, self.convolution): + register_page_pool_tensor(tensor=tensor) + + +class Qwen4ExpLexicalState(BoundedAttentionState): + def __init__(self, config, request_capacity, mtp_step=0, device="meta"): + self.config = config + self.state_name = None + self.request_capacity = request_capacity + self.mtp_step = mtp_step + count = len(config.logical_layer_ids) + self.ngram = torch.full( + (count, request_capacity, config.ngram_context), config.eos_token_id, dtype=torch.int64, device=device + ) + self.convolution = torch.zeros( + (count, request_capacity, config.conv_context, config.width), dtype=config.dtype, device=device + ) + self.big_pages_pool = self.small_pages_pool = None + self.lexical_caches = {} + self.speculative_lexical_caches = {} + + def __getstate__(self): + state = self.__dict__.copy() + state["lexical_caches"] = {} + state["speculative_lexical_caches"] = {} + return state + + def get_state_type(self): + return "qwen4_lexical" + + def contains_logical_layer(self, logical_layer_id): + return logical_layer_id in self.config.logical_layer_ids + + def get_layer_view(self, logical_layer_id): + layer = self.config.logical_layer_ids.index(logical_layer_id) + return self.ngram[layer], self.convolution[layer] + + def bind_weights(self, layer_weights): + for weight in layer_weights: + if not self.contains_logical_layer(weight.layer_num_): + continue + cache = LexicalRequestCache.__new__(LexicalRequestCache) + cache.layer_weight = weight + cache.eos_token_id = self.config.eos_token_id + cache.ngram, cache.convolution = self.get_layer_view(weight.layer_num_) + self.lexical_caches[weight.layer_num_] = cache + if self.mtp_step: + self.speculative_lexical_caches[weight.layer_num_] = LexicalSpeculativeCache(cache, self.mtp_step + 1) + + def prepare_lexical(self, metadata, reset_empty_prefix=False): + return {layer: cache.prepare(metadata, reset_empty_prefix) for layer, cache in self.lexical_caches.items()} + + def commit_lexical(self, metadata, batches, accepted_lengths=None): + for layer, cache in self.lexical_caches.items(): + cache.commit(metadata, batches[layer], accepted_lengths) + + def stage_speculative_lexical(self, layer, metadata, batch): + self.speculative_lexical_caches[layer].stage(metadata, batch) + + def initialize_page_pools(self, big_page_capacity, small_page_capacity=0, device="cpu", pin_memory=True): + self.big_pages_pool = ( + LexicalPages(big_page_capacity, self.config, device, pin_memory) if big_page_capacity else None + ) + self.small_pages_pool = ( + LexicalPages(small_page_capacity, self.config, device, pin_memory) if small_page_capacity else None + ) + + def reset_bounded_att_state(self, req_indices): + for request in req_indices: + self.ngram[:, request].fill_(self.config.eos_token_id) + self.convolution[:, request].zero_() + + def _pool(self, page_kind): + assert page_kind in ("big", "small") + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} lexical pages are disabled" + return pool + + def _copy_pages(self, req_indices, state_token_lens, page_kind, page_slot_ids, write): + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids) + pool = self._pool(page_kind) + for name in ("ngram", "convolution"): + copies = [ + (getattr(self, name)[:, req], getattr(pool, name)[slot]) + for req, slot in zip(req_indices, page_slot_ids) + ] + if not write: + copies = [(dst, src) for src, dst in copies] + copy_bounded_att_tensors(copies=copies) + + def save_bounded_att_runtime_to_pages(self, req_indices, state_token_lens, page_kind, page_slot_ids): + self._copy_pages(req_indices, state_token_lens, page_kind, page_slot_ids, True) + + def restore_bounded_att_runtime_from_pages(self, req_indices, state_token_lens, page_kind, page_slot_ids): + self._copy_pages(req_indices, state_token_lens, page_kind, page_slot_ids, False) + + def commit_bounded_att_state(self, req_indices, accepted_steps): + starts = torch.arange(req_indices.numel(), dtype=torch.int32, device=req_indices.device) + for cache in self.speculative_lexical_caches.values(): + cache.commit(req_indices, starts, accepted_steps, self.request_capacity - 1) + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self): + return self._regions() + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self): + return self._regions() + + def _regions(self): + regions = [] + offset = 0 + for name, tensor in (("ngram", self.ngram), ("convolution", self.convolution)): + shape = tensor.shape[:1] + tensor.shape[2:] + axes = ("layer", "context") if name == "ngram" else ("layer", "context", "dim") + end = offset + prod(shape) * tensor.element_size() + regions.append( + TransferRegion( + field=name, + byte_start=offset, + byte_end=end, + dtype=tensor.dtype, + axes=axes, + shape=tuple(shape), + axis_mappings=tuple( + AxisMapping(axis=axis, local_start=0, external_start=0, length=size) + for axis, size in zip(axes[1:], shape[1:]) + ), + write_enabled=self.config.tp_rank == 0, + ) + ) + offset = end + return tuple(regions) + + def _fields(self, ngram, convolution): + return { + "ngram": TensorField(ngram, ("layer", "context")), + "convolution": TensorField(convolution, ("layer", "context", "dim")), + } + + def _transfer(self, fields, buffers, pages, layout, write): + assert layout.state_name == self.state_name + for copies in bounded_att_buffer_copies(fields, buffers, layout, pages, write): + copy_bounded_att_tensors(copies=copies) + + def write_state_pages_to_cpu( + self, page_kind, page_slot_ids, cpu_page_buffers, cpu_page_indices, layout, cpu_page_ready_flags=None + ): + assert len(page_slot_ids) == len(cpu_page_indices) + if cpu_page_ready_flags is not None: + assert len(cpu_page_ready_flags) == len(cpu_page_indices) + pool = self._pool(page_kind) + selected = [i for i in range(len(page_slot_ids)) if cpu_page_ready_flags is None or not cpu_page_ready_flags[i]] + self._transfer( + [self._fields(pool.ngram[page_slot_ids[i]], pool.convolution[page_slot_ids[i]]) for i in selected], + cpu_page_buffers, + [cpu_page_indices[i] for i in selected], + layout, + True, + ) + + def read_state_pages_from_cpu(self, cpu_page_buffers, cpu_page_indices, big_page_slot_ids, layout): + assert len(cpu_page_indices) == len(big_page_slot_ids) + pool = self._pool("big") + self._transfer( + [self._fields(pool.ngram[slot], pool.convolution[slot]) for slot in big_page_slot_ids], + cpu_page_buffers, + cpu_page_indices, + layout, + False, + ) + + def copy_big_pages_from_dp(self, src_page_ids, dst_page_ids, remote_dp_bounded_state): + assert self.config == remote_dp_bounded_state.config + assert len(src_page_ids) == len(dst_page_ids) + for name in ("ngram", "convolution"): + copy_bounded_att_tensors( + copies=[ + (getattr(remote_dp_bounded_state._pool("big"), name)[src], getattr(self._pool("big"), name)[dst]) + for src, dst in zip(src_page_ids, dst_page_ids) + ] + ) + + def copy_runtime_state_from_dp(self, src_req_indices, dst_req_indices, remote_dp_bounded_state): + assert self.config == remote_dp_bounded_state.config + assert len(src_req_indices) == len(dst_req_indices) + for name in ("ngram", "convolution"): + copy_bounded_att_tensors( + copies=[ + (getattr(remote_dp_bounded_state, name)[:, src], getattr(self, name)[:, dst]) + for src, dst in zip(src_req_indices, dst_req_indices) + ] + ) + + def write_to_pd_page(self, req_idx, pd_page_buffers, pd_page_index, state_token_len, layout): + self._transfer( + [self._fields(self.ngram[:, req_idx], self.convolution[:, req_idx])], + pd_page_buffers, + [pd_page_index], + layout, + True, + ) + + def read_from_pd_page(self, pd_page_buffers, req_idx, pd_page_index, state_token_len, layout): + self._transfer( + [self._fields(self.ngram[:, req_idx], self.convolution[:, req_idx])], + pd_page_buffers, + [pd_page_index], + layout, + False, + ) diff --git a/lightllm/models/qwen4_exp/model.py b/lightllm/models/qwen4_exp/model.py new file mode 100644 index 0000000000..e5112afed1 --- /dev/null +++ b/lightllm/models/qwen4_exp/model.py @@ -0,0 +1,114 @@ +import json +from pathlib import Path + +from lightllm.models.qwen3_5_moe.model import Qwen3_5MOETpPartModel +from lightllm.models.registry import ModelRegistry +from lightllm.common.basemodel.attention.triton.fp import TritonAttBackend +from lightllm.utils.torch_dtype_utils import get_torch_dtype +from .configuration import normalize_model_config +from .infer_struct import Qwen4ExpInferStateInfo +from .layer_weights.pre_and_post_layer_weight import Qwen4ExpPreAndPostLayerWeight +from .layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight +from .layer_infer.pre_layer_infer import Qwen4ExpPreLayerInfer +from .layer_infer.post_layer_infer import Qwen4ExpPostLayerInfer +from .layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer + + +class Qwen4ExpTpPartModel(Qwen3_5MOETpPartModel): + linear_att_mtp_match_decode = True + pre_and_post_weight_class = Qwen4ExpPreAndPostLayerWeight + transformer_weight_class = Qwen4ExpTransformerLayerWeight + pre_layer_infer_class = Qwen4ExpPreLayerInfer + post_layer_infer_class = Qwen4ExpPostLayerInfer + transformer_layer_infer_class = Qwen4ExpTransformerLayerInfer + infer_state_class = Qwen4ExpInferStateInfo + + def _init_config(self): + with (Path(self.weight_dir_) / "config.json").open() as stream: + checkpoint = json.load(stream) + self.config, self.vision_config = normalize_model_config(checkpoint, self.finetune_config) + self.num_kv_heads = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) + + @classmethod + def normalize_model_config(cls, model_config): + config = normalize_model_config(model_config)[0] + interval = config["full_attention_interval"] + expected = [ + "full_attention" if (i + 1) % interval == 0 else "linear_attention" + for i in range(config["num_hidden_layers"]) + ] + actual = [ + "full_attention" if kind == "indexed_attention" else kind for kind in config.get("layer_types", expected) + ] + if actual != expected: + raise ValueError("Qwen4Exp requires the configured periodic attention layout") + config["layer_types"] = expected + return config + + @classmethod + def create_full_att_states(cls, model_config, args, model_logical_layer_ids, tp_rank=0, tp_world_size=1): + from lightllm.common.attention_state.full_att.gqa.gqa_state import GQAAttentionConfig + from .attention_state import Qwen4ExpAttentionState + + config = cls.normalize_model_config(model_config) + if args.llm_kv_type != "None": + raise ValueError("Qwen4Exp QSA currently requires unquantized KV storage") + interval = config["full_attention_interval"] + layers = tuple((layer,) for index, layer in enumerate(model_logical_layer_ids) if (index + 1) % interval == 0) + return [ + Qwen4ExpAttentionState( + GQAAttentionConfig( + logical_layer_groups=layers, + global_head_num=config["num_key_value_heads"], + head_dim=config["head_dim"], + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + config["indexer_head_dim"], + int(config["head_dim"] * config["partial_rotary_factor"]), + args.page_size, + ) + ] + + @classmethod + def create_bounded_att_states( + cls, model_config, args, model_logical_layer_ids, request_capacity=1, tp_rank=0, tp_world_size=1, device="meta" + ): + from .lexical_state import LexicalStateConfig, Qwen4ExpLexicalState + from .checkpoint import ngram_eos_token_id + + config = cls.normalize_model_config(model_config) + states = super().create_bounded_att_states( + config, args, model_logical_layer_ids, request_capacity, tp_rank, tp_world_size, device + ) + layers = tuple(model_logical_layer_ids[layer - 1] for layer in config.get("ple_layer_ids", [])) + if layers: + states.append( + Qwen4ExpLexicalState( + LexicalStateConfig( + logical_layer_ids=layers, + ngram_context=config["ngram_size"] - 1, + conv_context=(config["ple_conv_kernel_size"] - 1) * config["ngram_size"], + width=config["hc_count"] * config["hidden_size"], + eos_token_id=ngram_eos_token_id(config), + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + ), + request_capacity, + args.mtp_step, + device, + ) + ) + return states + + def prepare_attention_states(self): + super().prepare_attention_states() + from .lexical_state import Qwen4ExpLexicalState + + state = getattr(self.att_state_manager, f"{self.att_state_role}_qwen4_lexical_state", None) + if isinstance(state, Qwen4ExpLexicalState): + state.bind_weights(self.trans_layers_weight) + + def _init_att_backend(self): + self.prefill_att_backend = self.decode_att_backend = TritonAttBackend(model=self) diff --git a/lightllm/models/qwen4_exp/mtp/__init__.py b/lightllm/models/qwen4_exp/mtp/__init__.py new file mode 100644 index 0000000000..2ae28399f5 --- /dev/null +++ b/lightllm/models/qwen4_exp/mtp/__init__.py @@ -0,0 +1 @@ +pass diff --git a/lightllm/models/qwen4_exp/mtp/infer.py b/lightllm/models/qwen4_exp/mtp/infer.py new file mode 100644 index 0000000000..6e3ff96e3d --- /dev/null +++ b/lightllm/models/qwen4_exp/mtp/infer.py @@ -0,0 +1,29 @@ +import torch + +from lightllm.models.llama.layer_infer.pre_layer_infer import LlamaPreLayerInfer +from ..layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer +from ..multimodal import canonical_image_ids + + +class Qwen4ExpMTPPreLayerInfer(LlamaPreLayerInfer): + def __init__(self, network_config): + super().__init__(network_config) + self.image_token_id = network_config.get("image_token_id") + + def context_forward(self, input_ids, infer_state, layer_weight): + images = [image for params in (infer_state.multimodal_params or []) for image in params.get("images", [])] + if images: + starts = torch.tensor([image["token_id"] for image in images], dtype=torch.int64, device=input_ids.device) + lengths = torch.tensor([image["token_num"] for image in images], dtype=torch.int64, device=input_ids.device) + input_ids = canonical_image_ids(input_ids, starts, lengths, self.image_token_id) + return self.token_forward(input_ids, infer_state, layer_weight) + + def token_forward(self, input_ids, infer_state, layer_weight): + embedding = super().token_forward(input_ids, infer_state, layer_weight) + return layer_weight.fusion.component(embedding, infer_state.mtp_draft_input_hiddens) + + +class Qwen4ExpMTPTransformerLayerInfer(Qwen4ExpTransformerLayerInfer): + def __init__(self, layer_num, network_config): + config = dict(network_config, layer_types=None, full_attention_interval=1) + super().__init__(layer_num, config) diff --git a/lightllm/models/qwen4_exp/mtp/model.py b/lightllm/models/qwen4_exp/mtp/model.py new file mode 100644 index 0000000000..7c28fc649d --- /dev/null +++ b/lightllm/models/qwen4_exp/mtp/model.py @@ -0,0 +1,51 @@ +import torch + +from lightllm.models.draft_registry import DraftModelRegistry +from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel +from ..model import Qwen4ExpTpPartModel +from ..infer_struct import Qwen4ExpInferStateInfo +from ..layer_infer.post_layer_infer import Qwen4ExpPostLayerInfer +from .weights import Qwen4ExpMTPPreAndPostLayerWeight, Qwen4ExpMTPTransformerLayerWeight +from .infer import Qwen4ExpMTPPreLayerInfer, Qwen4ExpMTPTransformerLayerInfer + + +class Qwen4ExpMTPModel(Qwen3_5MTPModel): + pre_and_post_weight_class = Qwen4ExpMTPPreAndPostLayerWeight + transformer_weight_class = Qwen4ExpMTPTransformerLayerWeight + pre_layer_infer_class = Qwen4ExpMTPPreLayerInfer + post_layer_infer_class = Qwen4ExpPostLayerInfer + transformer_layer_infer_class = Qwen4ExpMTPTransformerLayerInfer + infer_state_class = Qwen4ExpInferStateInfo + _init_att_backend = Qwen4ExpTpPartModel._init_att_backend + create_full_att_states = classmethod(Qwen4ExpTpPartModel.create_full_att_states.__func__) + + @classmethod + def normalize_model_config(cls, model_config): + config = Qwen4ExpTpPartModel.normalize_model_config(model_config) + config.update( + full_attention_interval=1, num_hidden_layers=1, n_layer=1, ple_layer_ids=[], layer_types=["full_attention"] + ) + return config + + def _init_config(self): + Qwen4ExpTpPartModel._init_config(self) + if self.config.get("mtp_num_hidden_layers", 1) != 1: + raise ValueError("Qwen4Exp recurrent MTP requires one trained draft layer") + self.config.update( + full_attention_interval=1, num_hidden_layers=1, n_layer=1, ple_layer_ids=[], layer_types=["full_attention"] + ) + + def _init_weights(self, start_layer_index=None): + if start_layer_index is not None or self.mtp_previous_draft_models: + raise ValueError("Qwen4Exp MTP uses one recurrent predictor") + self.pre_post_weight = self.pre_and_post_weight_class(self.data_type, self.config) + self.pre_post_weight.wte_weight_ = self.main_model.pre_post_weight.wte_weight_ + self.pre_post_weight.lm_head_weight_ = self.main_model.pre_post_weight.lm_head_weight_ + self.trans_layers_weight = [self.transformer_weight_class(0, self.data_type, self.config, self.quant_cfg)] + + def _gen_special_model_input(self, token_num): + return { + "mtp_draft_input_hiddens": torch.randn( + token_num, self.config["hidden_size"] * self.config["hc_count"], dtype=self.data_type, device="cuda" + ) + } diff --git a/lightllm/models/qwen4_exp/mtp/weights.py b/lightllm/models/qwen4_exp/mtp/weights.py new file mode 100644 index 0000000000..d0edb3e5ea --- /dev/null +++ b/lightllm/models/qwen4_exp/mtp/weights.py @@ -0,0 +1,55 @@ +from dataclasses import dataclass +from functools import partial + +import torch.nn.functional as F + +from lightllm.common.basemodel import PreAndPostLayerWeight +from lightllm.common.gated_residual.kernels import grouped_norm +from ..layer_weights.components import ComponentWeight, gated_residual_weight +from ..layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight + + +@dataclass(frozen=True) +class MTPFusion: + parameters: dict + branches: int + eps: float + + def __call__(self, embedding, hidden): + if embedding.ndim != 2 or not embedding.is_contiguous() or not hidden.is_contiguous(): + raise ValueError("Qwen MTP fusion requires contiguous embedding and hidden matrices") + width = embedding.shape[-1] + if hidden.shape != (embedding.shape[0], self.branches * width): + raise ValueError("Qwen MTP requires the full pre-mixer HC hidden stream") + weights = self.parameters + embedding = grouped_norm(embedding, weights["mtp.pre_fc_norm_embedding.weight"], 1, self.eps) + hidden = grouped_norm(hidden, weights["mtp.pre_fc_norm_hidden.weight"], 1, self.eps) + embedding = F.linear(embedding, weights["mtp.fc_embedding.weight"]) + hidden = F.linear(hidden.view(-1, self.branches, width), weights["mtp.fc_hidden.weight"]) + return (hidden + embedding[:, None]).flatten(1) + + +class Qwen4ExpMTPPreAndPostLayerWeight(PreAndPostLayerWeight): + def __init__(self, data_type, network_config): + super().__init__(data_type, network_config) + hidden, branches = network_config["hidden_size"], network_config["hc_count"] + specs = { + "mtp.pre_fc_norm_embedding.weight": ((hidden,), data_type), + "mtp.pre_fc_norm_hidden.weight": ((hidden * branches,), data_type), + "mtp.fc_embedding.weight": ((hidden, hidden), data_type), + "mtp.fc_hidden.weight": ((hidden, hidden), data_type), + } + self.fusion = ComponentWeight(specs, partial(MTPFusion, branches=branches, eps=network_config["rms_norm_eps"])) + self.hyper_connection_mixer = gated_residual_weight( + "mtp.hyper_connection_mixer", network_config, data_type, use_combine=False + ) + self.wte_weight_ = self.lm_head_weight_ = None + + +class Qwen4ExpMTPTransformerLayerWeight(Qwen4ExpTransformerLayerWeight): + def load_hf_weights(self, weights): + prefix = "mtp.layers.0." + mapped = { + "model.layers.0." + name[len(prefix) :]: value for name, value in weights.items() if name.startswith(prefix) + } + super().load_hf_weights(mapped) diff --git a/lightllm/models/qwen4_exp/multimodal.py b/lightllm/models/qwen4_exp/multimodal.py new file mode 100644 index 0000000000..3978fd502a --- /dev/null +++ b/lightllm/models/qwen4_exp/multimodal.py @@ -0,0 +1,43 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _canonical_image_ids( + Tokens, + Starts, + Lengths, + Output, + COUNT: tl.constexpr, + IMAGES: tl.constexpr, + IMAGE_TOKEN: tl.constexpr, + BLOCK: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + values = tl.load(Tokens + offsets, offsets < COUNT, other=0).to(tl.int64) + canonical = values + for image in range(IMAGES): + start = tl.load(Starts + image).to(tl.int64) + length = tl.load(Lengths + image).to(tl.int64) + canonical = tl.where((values >= start) & (values < start + length), IMAGE_TOKEN, canonical) + tl.store(Output + offsets, canonical, offsets < COUNT) + + +def canonical_image_ids(tokens, starts, lengths, image_token_id): + if image_token_id is None: + raise ValueError("Qwen4 image hashing requires the checkpoint image token ID") + for value in (tokens, starts, lengths): + if value.ndim != 1 or not value.is_cuda or not value.is_contiguous() or value.device != tokens.device: + raise ValueError("Image token mapping requires contiguous CUDA vectors on one device") + if value.dtype not in (torch.int32, torch.int64): + raise ValueError("Image token mapping requires integer IDs and ranges") + if starts.shape != lengths.shape: + raise ValueError("Image cache starts and lengths must have matching shapes") + if not starts.numel() or not tokens.numel(): + return tokens + output = torch.empty_like(tokens) + _canonical_image_ids[(triton.cdiv(tokens.numel(), 256),)]( + tokens, starts, lengths, output, tokens.numel(), starts.numel(), image_token_id, 256 + ) + return output diff --git a/lightllm/models/qwen4_exp/pinned_memory.py b/lightllm/models/qwen4_exp/pinned_memory.py new file mode 100644 index 0000000000..cfa4f14269 --- /dev/null +++ b/lightllm/models/qwen4_exp/pinned_memory.py @@ -0,0 +1,72 @@ +import math +import mmap + +import numpy as np +import torch + +# Keep the number of CUDA host-registration ranges bounded when DP replicates +# the offloaded table on every rank. The registration unit still aligns each +# range to a complete recurrent snapshot, so a larger chunk cannot split a DMA +# snapshot across registrations. +_CHUNK_BYTES = 1024 * 1024 * 1024 +_LARGE_BYTES = 512 * 1024 * 1024 + + +class _RegisteredArray(np.ndarray): + def __array_finalize__(self, source): + self.allocation = getattr(source, "allocation", None) + + +class _Registration: + def __init__(self, mapping, device, runtime): + self.mapping = mapping + self.device = device + self.runtime = runtime + self.addresses = [] + + def __del__(self): + if not self.addresses: + return + + try: + with torch.cuda.device(self.device): + torch.cuda.synchronize() + except Exception: + pass + for address in reversed(self.addresses): + try: + self.runtime.cudaHostUnregister(address) + except Exception: + pass + + +def pinned_zeros(shape, *, dtype, device="cpu", pin_memory=True): + size = math.prod(shape) * dtype.itemsize + if device != "cpu" or not pin_memory: + raise ValueError("Qwen recurrent snapshots require pinned CPU storage") + if size < _LARGE_BYTES: + return torch.zeros(shape, dtype=dtype, device="cpu", pin_memory=True) + + cell_bytes = math.prod(shape[1:]) * dtype.itemsize + registration_unit = math.lcm(cell_bytes, mmap.PAGESIZE) + if registration_unit > _CHUNK_BYTES: + raise ValueError("A recurrent snapshot must fit within one bounded host registration") + chunk_bytes = (_CHUNK_BYTES // registration_unit) * registration_unit + runtime = torch.cuda.cudart() + mapping = mmap.mmap(-1, size) + allocation = _Registration(mapping, torch.cuda.current_device(), runtime) + array = np.frombuffer(mapping, dtype=np.uint8).view(_RegisteredArray) + array.allocation = allocation + value = torch.from_numpy(array).view(dtype).reshape(shape) + value.zero_() + for offset in range(0, size, chunk_bytes): + address = value.data_ptr() + offset + status = runtime.cudaHostRegister(address, min(chunk_bytes, size - offset), 0) + if int(status): + raise RuntimeError( + f"Cannot register Qwen pinned snapshot at byte {offset}: {runtime.cudaGetErrorString(status)}" + ) + allocation.addresses.append(address) + + value._qwen4_pinned_allocation = allocation + return value diff --git a/lightllm/models/qwen4_exp/position.py b/lightllm/models/qwen4_exp/position.py new file mode 100644 index 0000000000..e5fcaa054a --- /dev/null +++ b/lightllm/models/qwen4_exp/position.py @@ -0,0 +1,18 @@ +import torch + + +def combine_mrope(cos: torch.Tensor, sin: torch.Tensor, sections): + if cos.shape != sin.shape or cos.device != sin.device or cos.dtype != sin.dtype: + raise ValueError("Rotary cosine and sine shapes, devices and dtypes must match") + if cos.ndim == 2: + return cos.repeat(1, 2).contiguous(), sin.repeat(1, 2).contiguous() + if cos.ndim != 3 or cos.shape[0] != 3 or len(sections) != 3: + raise ValueError("Expected three rotary axes and three mRoPE sections") + offsets = torch.arange(cos.shape[-1], device=cos.device) + height = ((offsets % 3) == 1) & (offsets <= 3 * sections[1]) + width = ((offsets % 3) == 2) & (offsets <= 3 * sections[2]) + combined = [] + for tensor in (cos, sin): + values = torch.where(height, tensor[1], torch.where(width, tensor[2], tensor[0])) + combined.append(values.repeat(1, 2).contiguous()) + return tuple(combined) diff --git a/lightllm/models/qwen4_exp/speculative_cache.py b/lightllm/models/qwen4_exp/speculative_cache.py new file mode 100644 index 0000000000..e1c51570d1 --- /dev/null +++ b/lightllm/models/qwen4_exp/speculative_cache.py @@ -0,0 +1,108 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _stage( + History, + Inputs, + Starts, + Requests, + Pending, + CONTEXT: tl.constexpr, + WIDTH: tl.constexpr, + VERIFY: tl.constexpr, + BLOCK: tl.constexpr, +): + segment = tl.program_id(0) + request = tl.load(Requests + segment).to(tl.int64) + if request >= 0: + start = tl.load(Starts + segment) + count = tl.load(Starts + segment + 1) - start + offset = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + position, column = offset // WIDTH, offset % WIDTH + old = tl.load(History + segment.to(tl.int64) * CONTEXT * WIDTH + offset, position < CONTEXT, other=0) + new = tl.load( + Inputs + (start.to(tl.int64) + position - CONTEXT) * WIDTH + column, + (position >= CONTEXT) & (position < CONTEXT + count), + other=0, + ) + tl.store( + Pending + request * (CONTEXT + VERIFY) * WIDTH + offset, + tl.where(position < CONTEXT, old, new), + (position < CONTEXT + count) & (position < CONTEXT + VERIFY), + ) + + +@triton.jit +def _commit( + Pending, + History, + Requests, + Starts, + Accepted, + HOLD: tl.constexpr, + CONTEXT: tl.constexpr, + WIDTH: tl.constexpr, + VERIFY: tl.constexpr, + BLOCK: tl.constexpr, +): + group = tl.program_id(0) + start = tl.load(Starts + group) + request = tl.load(Requests + start).to(tl.int64) + if request != HOLD: + count = tl.load(Accepted + group) + offset = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + values = tl.load( + Pending + request * (CONTEXT + VERIFY) * WIDTH + count * WIDTH + offset, offset < CONTEXT * WIDTH, other=0 + ) + tl.store(History + request * CONTEXT * WIDTH + offset, values, offset < CONTEXT * WIDTH) + + +class LexicalSpeculativeCache: + def __init__(self, accepted_cache, verify_width): + if verify_width < 1: + raise ValueError("Speculative lexical cache requires positive verify width") + self.verify_width = verify_width + self.fields = [] + for history in (accepted_cache.ngram, accepted_cache.convolution): + requests, context = history.shape[:2] + width = history.shape[2] if history.ndim == 3 else 1 + pending = torch.empty((requests, context + verify_width, width), dtype=history.dtype, device=history.device) + self.fields.append((history, pending, context, width)) + + def stage(self, metadata, batch): + if batch.input_tokens is None or batch.conv_inputs is None: + raise RuntimeError("PLE candidates must be computed before staging") + for (history, pending, context, width), prefix, inputs in zip( + self.fields, (batch.ngram_history, batch.conv_history), (batch.input_tokens, batch.conv_inputs) + ): + if metadata.active_rows.numel(): + _stage[(metadata.active_rows.numel(), triton.cdiv((context + self.verify_width) * width, 256))]( + prefix, + inputs, + metadata.starts, + metadata.active_rows, + pending, + context, + width, + self.verify_width, + 256, + ) + + def commit(self, requests, starts, accepted_lengths, hold_request): + for history, pending, context, width in self.fields: + if starts.numel() and context: + _commit[(starts.numel(), triton.cdiv(context * width, 256))]( + pending, + history, + requests, + starts, + accepted_lengths, + hold_request, + context, + width, + self.verify_width, + 256, + ) diff --git a/lightllm/models/qwen4_exp/state.py b/lightllm/models/qwen4_exp/state.py new file mode 100644 index 0000000000..6800f8a6e6 --- /dev/null +++ b/lightllm/models/qwen4_exp/state.py @@ -0,0 +1,136 @@ +from dataclasses import dataclass +from typing import Optional + +import torch +import triton + +from lightllm.common.qsa.attention import indexed_attention +from lightllm.distributed import all_reduce + + +@dataclass +class QSABatch: + slot_table: torch.Tensor + active_rows: torch.Tensor + key_lengths: torch.Tensor + first_blocks: torch.Tensor + query_rows: torch.Tensor + visible_lengths: torch.Tensor + mem_indices: torch.Tensor + cos: torch.Tensor + sin: torch.Tensor + max_new_blocks: int + + @classmethod + def from_metadata(cls, metadata, slot_table, mem_indices, cos, sin, compress_ratio: int): + if compress_ratio < 1: + raise ValueError("QSA compression ratio must be positive") + return cls( + slot_table, + metadata.active_rows, + metadata.key_lengths, + metadata.first_blocks, + metadata.query_rows, + metadata.visible_lengths, + mem_indices.to(torch.int64), + cos, + sin, + (mem_indices.numel() + compress_ratio - 1) // compress_ratio, + ) + + +@dataclass +class QSAExecutionState: + raw_keys: torch.Tensor + + pooled_keys: torch.Tensor + key_cos: torch.Tensor + key_sin: torch.Tensor + + def attend( + self, + indexer, + hidden, + query, + key, + value, + batch: QSABatch, + *, + decode: bool = False, + attention_splits: Optional[int] = None, + ): + index_query, raw = indexer.project(hidden, batch.cos, batch.sin) + self.raw_keys.index_copy_(0, batch.mem_indices, raw) + indexer.pool( + self.raw_keys, + batch.slot_table, + batch.key_lengths, + self.key_cos, + self.key_sin, + self.pooled_keys, + batch.first_blocks, + batch.max_new_blocks, + request_rows=batch.active_rows, + ) + selected = indexer.select( + index_query, + self.pooled_keys, + batch.query_rows, + batch.visible_lengths, + batch.slot_table, + query_chunk_size=max(64, query.shape[0]) if decode else 1024, + ) + + splits = attention_splits + if splits is None and decode: + splits = min(32, triton.cdiv(selected.shape[1], 64)) + wide_block = query.dtype in (torch.float16, torch.bfloat16) + return indexed_attention( + query.contiguous(), + key, + value, + selected, + splits=splits, + selection_lengths=batch.visible_lengths, + short_rows_direct=decode or attention_splits is not None, + skip_empty_splits=True, + block_size=128 if wide_block else 64, + num_stages=2 if wide_block else None, + ) + + +@dataclass +class PLEBatch: + starts: torch.Tensor + ngram_history: torch.Tensor + conv_history: torch.Tensor + mask: Optional[torch.Tensor] = None + conv_inputs: Optional[torch.Tensor] = None + input_tokens: Optional[torch.Tensor] = None + + def apply(self, residual, tokens, layer_weight, tp_world_size: int, dist_group): + hasher = layer_weight.ngram_hasher.component + if self.mask is not None: + tokens = tokens.masked_fill(~self.mask, hasher.eos_token_id) + indices, _ = hasher.hash(tokens, self.starts, self.ngram_history) + embeddings = layer_weight.ngram_embedding(indices.flatten()).view( + tokens.numel(), hasher.heads * layer_weight.ngram_embedding.dim + ) + if tp_world_size > 1: + all_reduce(embeddings, group=dist_group, async_op=False) + update, conv_inputs = layer_weight.ple.component.forward( + residual, embeddings, self.starts, self.conv_history, self.mask + ) + self.input_tokens, self.conv_inputs = tokens, conv_inputs + return update + + def histories_after(self, layer_weight, accepted_lengths=None): + if self.input_tokens is None or self.conv_inputs is None: + raise RuntimeError("PLE candidates must be computed before requesting histories") + ngram = layer_weight.ngram_hasher.component.history_after( + self.input_tokens, self.starts, self.ngram_history, accepted_lengths + ) + conv = layer_weight.ple.component.history_after( + self.conv_inputs, self.starts, self.conv_history, accepted_lengths + ) + return ngram, conv diff --git a/lightllm/models/qwen4_exp/tokenizer.py b/lightllm/models/qwen4_exp/tokenizer.py new file mode 100644 index 0000000000..807cea9ad5 --- /dev/null +++ b/lightllm/models/qwen4_exp/tokenizer.py @@ -0,0 +1,6 @@ +from lightllm.models.qwen2_vl.model import QWen2VLTokenizer + + +class Qwen4ExpTokenizer(QWen2VLTokenizer): + def count_prompt_tokens(self, text_tokens, image_tokens, image_count, audio_tokens, audio_count): + return text_tokens + image_tokens + audio_tokens diff --git a/lightllm/models/qwen4_exp/triton_kernel/__init__.py b/lightllm/models/qwen4_exp/triton_kernel/__init__.py new file mode 100644 index 0000000000..2ae28399f5 --- /dev/null +++ b/lightllm/models/qwen4_exp/triton_kernel/__init__.py @@ -0,0 +1 @@ +pass diff --git a/lightllm/models/qwen4_exp/triton_kernel/cache_pages.py b/lightllm/models/qwen4_exp/triton_kernel/cache_pages.py new file mode 100644 index 0000000000..c541c1b892 --- /dev/null +++ b/lightllm/models/qwen4_exp/triton_kernel/cache_pages.py @@ -0,0 +1,71 @@ +import triton +import triton.language as tl + + +@triton.jit +def copy_slot_pages( + Slots, + Pages, + MemIndices, + PageIndices, + Readies, + SLOT_STRIDE: tl.constexpr, + LAYER_STRIDE: tl.constexpr, + PAGE_STRIDE: tl.constexpr, + TOKENS: tl.constexpr, + LAYERS: tl.constexpr, + DIM: tl.constexpr, + OFFLOAD: tl.constexpr, + BLOCK: tl.constexpr, +): + page = tl.program_id(0) + page_index = tl.load(PageIndices + page).to(tl.int64) + active = page_index >= 0 + if OFFLOAD: + active = active & ~tl.load(Readies + page).to(tl.int1) + if active: + col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = col < TOKENS * LAYERS * DIM + token = col // (LAYERS * DIM) + layer = ((col // DIM) % LAYERS).to(tl.int64) + slot = tl.load(MemIndices + page * TOKENS + token, mask, other=-1).to(tl.int64) + mask = mask & (slot >= 0) + live = Slots + slot * SLOT_STRIDE + layer * LAYER_STRIDE + col % DIM + cached = Pages + page_index * PAGE_STRIDE + col + if OFFLOAD: + value = tl.load(live, mask, other=0) + tl.store(cached, value, mask) + else: + value = tl.load(cached, mask, other=0) + tl.store(live, value, mask) + + +@triton.jit +def copy_state_pages( + States, + Pages, + StateIndices, + PageIndices, + Readies, + WIDTH: tl.constexpr, + PAGE_STRIDE: tl.constexpr, + OFFLOAD: tl.constexpr, + BLOCK: tl.constexpr, +): + page = tl.program_id(0) + page_index = tl.load(PageIndices + page).to(tl.int64) + active = page_index >= 0 + if OFFLOAD: + active = active & ~tl.load(Readies + page).to(tl.int1) + if active: + state = tl.load(StateIndices + page).to(tl.int64) + col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = col < WIDTH + live = States + state * WIDTH + col + cached = Pages + page_index * PAGE_STRIDE + col + if OFFLOAD: + value = tl.load(live, mask, other=0) + tl.store(cached, value, mask) + else: + value = tl.load(cached, mask, other=0) + tl.store(live, value, mask) diff --git a/lightllm/models/qwen4_exp/triton_kernel/gated_norm.py b/lightllm/models/qwen4_exp/triton_kernel/gated_norm.py new file mode 100644 index 0000000000..9f50ce3acc --- /dev/null +++ b/lightllm/models/qwen4_exp/triton_kernel/gated_norm.py @@ -0,0 +1,92 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _gated_norm( + X, + Gate, + Weight, + Out, + input_stride_row, + gate_stride_token, + gate_stride_head, + out_stride_row, + GATE_HEADS: tl.constexpr, + D: tl.constexpr, + EPS: tl.constexpr, + SIGMOID: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + col = tl.arange(0, BLOCK) + x = tl.load(X + row * input_stride_row + col, col < D, other=0).to(tl.float32) + inv = tl.rsqrt(tl.sum(x * x, 0) / D + EPS) + + normalized = (x * inv).to(X.dtype.element_ty).to(tl.float32) + weight = tl.load(Weight + col, col < D, other=0).to(tl.float32) + weighted = (normalized * weight).to(X.dtype.element_ty).to(tl.float32) + gate_token = row // GATE_HEADS + gate_head = row % GATE_HEADS + gate = tl.load( + Gate + gate_token * gate_stride_token + gate_head * gate_stride_head + col, + col < D, + other=0, + ).to(tl.float32) + activation = tl.sigmoid(gate) + if not SIGMOID: + activation = gate * activation + tl.store(Out + row * out_stride_row + col, weighted * activation, col < D) + + +@torch.no_grad() +def gated_norm(input: torch.Tensor, gate: torch.Tensor, weight: torch.Tensor, eps: float, activation: str, out=None): + if activation not in ("sigmoid", "silu"): + raise ValueError("Qwen4Exp GDN output gate must be sigmoid or silu") + if ( + input.ndim != 2 + or input.shape[-1] != weight.numel() + or gate.ndim not in (2, 3) + or gate.shape[-1] != input.shape[-1] + or (gate.ndim == 2 and gate.shape != input.shape) + or (gate.ndim == 3 and gate.shape[0] * gate.shape[1] != input.shape[0]) + or not input.is_contiguous() + or gate.stride(-1) != 1 + ): + raise ValueError("GDN norm requires a contiguous input, a packed gate, and a per-head norm weight") + for tensor in (input, gate, weight): + if not tensor.is_cuda or tensor.dtype != input.dtype or tensor.device != input.device: + raise ValueError("GDN norm tensors must be CUDA tensors with matching device and dtype") + if weight.ndim != 1 or not weight.is_contiguous() or input.ndim < 2 or eps <= 0: + raise ValueError("GDN norm requires a weight vector, token/head rows and positive epsilon") + result = torch.empty_like(input) if out is None else out + if ( + result.shape != input.shape + or result.dtype != input.dtype + or result.device != input.device + or not result.is_contiguous() + ): + raise ValueError("GDN norm output must be contiguous and match the input") + rows = input.shape[0] + gate_heads = 1 if gate.ndim == 2 else gate.shape[1] + gate_stride_token = gate.stride(0) + gate_stride_head = 0 if gate.ndim == 2 else gate.stride(1) + if rows: + _gated_norm[(rows,)]( + input, + gate, + weight, + result, + input.stride(0), + gate_stride_token, + gate_stride_head, + result.stride(0), + GATE_HEADS=gate_heads, + D=input.shape[-1], + EPS=eps, + SIGMOID=activation == "sigmoid", + BLOCK=triton.next_power_of_2(input.shape[-1]), + enable_fp_fusion=False, + ) + return result diff --git a/lightllm/models/qwen4_exp/vision.py b/lightllm/models/qwen4_exp/vision.py new file mode 100644 index 0000000000..93b1d49c19 --- /dev/null +++ b/lightllm/models/qwen4_exp/vision.py @@ -0,0 +1,36 @@ +import torch +import torch.nn.functional as F + +from lightllm.models.qwen3_vl.qwen3_visual import Qwen3VisionTransformerPretrainedModel +from lightllm.models.qwen2_vl.qwen2_visual import VisionRotaryEmbedding + + +class Qwen4ExpVisionTransformer(Qwen3VisionTransformerPretrainedModel): + def __init__(self, kvargs, **config): + config = dict(config) + if config.get("deepstack_visual_indexes"): + raise ValueError("Qwen4Exp vision does not define deepstack outputs") + config["deepstack_visual_indexes"] = [] + rope = config.get("rope_parameters") or {"rope_type": "axial", "rope_theta": 10000.0} + if rope.get("rope_type", "axial") != "axial": + raise ValueError("Qwen4Exp vision requires axial rotary embeddings") + super().__init__(kvargs, **config) + self.rotary_pos_emb = VisionRotaryEmbedding( + self.hidden_size // self.num_heads // 2, theta=rope.get("rope_theta", 10000.0) + ).cuda() + + def fast_pos_embed_interpolate(self, grid_thw): + table = self.pos_embed.weight.float().T.reshape( + 1, self.hidden_size, self.num_grid_per_side, self.num_grid_per_side + ) + merge = self.spatial_merge_size + patches = [] + for frames, height, width in grid_thw.detach().cpu().tolist(): + if min(frames, height, width) < 1 or height % merge or width % merge: + raise ValueError("Vision grids must be positive and divisible by the spatial merge size") + positions = F.interpolate(table, size=(height, width), mode="bilinear", align_corners=True) + positions = positions[0].permute(1, 2, 0) + positions = positions.reshape(height // merge, merge, width // merge, merge, self.hidden_size) + positions = positions.permute(0, 2, 1, 3, 4).reshape(-1, self.hidden_size).repeat(frames, 1) + patches.append(positions.to(self.pos_embed.weight.dtype)) + return torch.cat(patches) diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 41ae769ae6..4bdfbdbf74 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -1,4 +1,5 @@ import argparse +import json def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: @@ -945,6 +946,11 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: help="""enable torch memory saver, which is used for release_memory and resume_memory during RL training.""", ) parser.add_argument("--enable_weight_cpu_backup", action="store_true", help="""enable weight cpu backup.""") + parser.add_argument( + "--offload_ngram_embedding", + action="store_true", + help="Keep supported n-gram embedding tables in pinned CPU memory; gather requested rows onto CUDA.", + ) parser.add_argument( "--disk_cache_dir", type=str, @@ -984,6 +990,21 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: If radix cache hit rate is low under high load due to limited small-page capacity and LRU eviction, increasing linear_att_cache_size can improve hit rate at the cost of more memory.""", ) + parser.add_argument( + "--ssm_state_mode", + choices=["native", "replay"], + default="native", + help="SSM state storage: native snapshots or deferred GDN/KDA ReplaySSM (FP32/BF16).", + ) + parser.add_argument( + "--ssm_replay_config", + type=json.loads, + default={}, + help='Replay-only JSON object: {"cache_len":16,"compute_mode":"precompute"}. ' + "cache_len: 4/8/16/32/64, at least the verify width; compute_mode: inline/precompute. " + "mtp_step=0 decode uses the same replay history: explicit precompute projects that one token, " + "and the default is inline. MTP defaults to precompute.", + ) parser.add_argument( "--linear_att_ssm_data_type", type=str, diff --git a/lightllm/server/api_lightllm.py b/lightllm/server/api_lightllm.py index b881b8ad1f..1fc24a127b 100644 --- a/lightllm/server/api_lightllm.py +++ b/lightllm/server/api_lightllm.py @@ -30,7 +30,6 @@ async def lightllm_get_score(request: Request, httpserver_manager: HttpServerMan async def lightllm_generate(request: Request, httpserver_manager: HttpServerManager) -> Response: - request_dict = await request.json() prompt = request_dict.pop("inputs") sample_params_dict = request_dict["parameters"] @@ -115,11 +114,10 @@ async def lightllm_generate(request: Request, httpserver_manager: HttpServerMana async def lightllm_generate_stream(request: Request, httpserver_manager: HttpServerManager) -> Response: - request_dict = await request.json() prompt = request_dict.pop("inputs") sample_params_dict = request_dict["parameters"] - _ = sample_params_dict.pop("return_details", False) + return_details = sample_params_dict.pop("return_details", False) _ = sample_params_dict.pop("return_routed_experts", None) sampling_params = SamplingParams() sampling_params.init(tokenizer=httpserver_manager.tokenizer, **sample_params_dict) @@ -155,6 +153,15 @@ async def stream_results() -> AsyncGenerator[bytes, None]: "input_usage": input_usage, } ret["token"]["logprobs"] = metadata["logprobs"] + if return_details: + for name in ( + "prompt_cache_len", + "mtp_accepted_token_num", + "mtp_verify_token_num", + "mtp_verify_step_num", + ): + if name in metadata: + ret["token"][name] = metadata[name] if "prompt_logprobs" in metadata: ret["prompt_logprobs"] = metadata["prompt_logprobs"] ret["prompt_token_ids"] = metadata.get("prompt_token_ids") diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index a42acef884..361d789a64 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -43,6 +43,59 @@ def _set_envs_and_config(args: StartArgs): mp.set_start_method("spawn", force=True) +def _validate_and_normalize_ssm_config(args: StartArgs): + assert args.ssm_state_mode in ("native", "replay"), "Invalid ssm_state_mode" + config = args.ssm_replay_config + assert isinstance(config, dict), "ssm_replay_config must be a JSON object" + assert args.ssm_state_mode == "replay" or not config, "ssm_replay_config requires ssm_state_mode=replay" + if args.ssm_state_mode == "replay": + defaults = {"cache_len": 16, "compute_mode": "precompute" if args.mtp_step > 0 else "inline"} + unknown = config.keys() - defaults.keys() + assert not unknown, f"Unknown ssm_replay_config fields: {unknown}" + config = {**defaults, **config} + assert type(config["cache_len"]) is int and config["cache_len"] in ( + 4, + 8, + 16, + 32, + 64, + ), "ReplaySSM capacity must be a power of two and at least 4, up to 64" + assert config["compute_mode"] in ("inline", "precompute"), "Invalid replay compute_mode" + args.ssm_replay_config = config + logger.info(f"SSM mode: replay, config: {config}") + else: + logger.info(f"SSM mode: {args.ssm_state_mode}") + from lightllm.utils.config_utils import get_model_type + + model_type = get_model_type(args.model_dir) + is_gdn_or_kda = model_type in ( + "glm5_next", + "glm5_next_text", + "qwen3_next", + "qwen3_5", + "qwen3_5_moe", + "qwen3_5_text", + "qwen3_5_moe_text", + "qwen4_exp", + "qwen4_exp_text", + "sensenova_loop", + ) + if is_gdn_or_kda: + assert args.mtp_step <= 15, ( + "GDN/KDA requires mtp_step <= 15: " "linear attention MTP state acceptance supports verify width <= 16" + ) + if args.ssm_state_mode == "native": + return + assert args.linear_att_ssm_data_type in ( + "float32", + "bfloat16", + ), "Non-native SSM state modes require FP32 or BF16 state" + assert is_gdn_or_kda, "Non-native SSM state modes currently require a GDN or KDA model" + assert args.mtp_step >= 0, "ssm_state_mode=replay requires mtp_step >= 0" + capacity = config["cache_len"] + assert capacity >= args.mtp_step + 1, "ReplaySSM capacity must cover the verify width" + + def _launch_subprocesses(args: StartArgs): _validate_pd_features(args) _set_envs_and_config(args) @@ -355,6 +408,8 @@ def _launch_subprocesses(args: StartArgs): s3_model_prepare(args.model_dir) + _validate_and_normalize_ssm_config(args) + # 如果args.eos_id 是 None, 从 config.json 中读取 eos_token_id 相关的信息，赋值给 args if args.eos_id is None: from lightllm.utils.config_utils import get_eos_token_ids diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 0d6468f397..0f43364d73 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -246,6 +246,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) + offload_ngram_embedding: bool = field(default=False) hardware_platform: str = field(default="cuda", metadata={"choices": ["cuda", "musa"]}) enable_torch_fallback: bool = field(default=False) enable_triton_fallback: bool = field(default=False) @@ -255,6 +256,8 @@ class StartArgs: weight_version: str = "default" # hybrid attention model (Qwen3Next) + ssm_state_mode: str = "native" + ssm_replay_config: dict[str, object] = field(default_factory=dict) linear_att_hash_page_size: int = field(default=512) linear_att_page_block_num: int = field(default=10000000) disable_linear_att_small_page_cpu_cache: bool = field(default=False) diff --git a/lightllm/server/httpserver/manager.py b/lightllm/server/httpserver/manager.py index 43b06f02ef..541e5b9502 100644 --- a/lightllm/server/httpserver/manager.py +++ b/lightllm/server/httpserver/manager.py @@ -1,3 +1,4 @@ +from lightllm.server.httpserver.slot_availability import RequestSlotAvailability import sys import zmq import zmq.asyncio @@ -56,6 +57,7 @@ def __init__( args: StartArgs, ): self.args: StartArgs = args + self._request_slot_availability = RequestSlotAvailability.from_args(args) ports = get_shm_port_args() context = zmq.asyncio.Context(2) self.send_to_router = context.socket(zmq.PUSH) @@ -344,7 +346,6 @@ async def _generate( # 用于等待 pd_master 下发的交换信息 pd_event: asyncio.Event = None, ) -> AsyncGenerator[Tuple[int, str, dict, FinishStatus], None]: - start_time = time.time() # 在请求登记、编码和 PD 握手前获取语法；缓存未命中时先完成编译，避免后续流程等待或校验失败。 compiled_grammar = await sampling_params.get_compiled_grammar(self.output_grammar_compiler) @@ -587,6 +588,8 @@ async def _alloc_shm_req_indexes( alloc_deadline = time.monotonic() + alloc_timeout_seconds if alloc_timeout_seconds is not None else None # 空闲槽不足时整组重试，避免多个 n > 1 的请求各占部分槽后互相等待。 + slot_availability = getattr(self, "_request_slot_availability", None) + observed_generation = slot_availability.generation if slot_availability is not None else 0 alloced_req_indexes = await self.shm_req_manager.async_alloc_req_indexes(req_num) # 保持相同的退避起点，仅通过系数让高优先级请求更快地重新尝试获取 shm_req。 sleep_time_factor = 0.2 if high_priority_request else 1 @@ -601,11 +604,20 @@ async def _alloc_shm_req_indexes( f"PD {self.args.run_mode} node is busy: unable to allocate a shm_req object " f"within {alloc_timeout_seconds} seconds" ) - await asyncio.sleep(sleep_time * sleep_time_factor) + if slot_availability is None: + await asyncio.sleep(sleep_time * sleep_time_factor) + else: + await slot_availability.wait(observed_generation, sleep_time * sleep_time_factor) + observed_generation = slot_availability.generation if slot_availability is not None else 0 sleep_time = min(1, sleep_time * 1.1) alloced_req_indexes = await self.shm_req_manager.async_alloc_req_indexes(req_num) return alloced_req_indexes + def _notify_req_slot_release(self): + availability = getattr(self, "_request_slot_availability", None) + if availability is not None: + availability.notify_release() + async def _log_req_header(self, request_headers, group_request_id: int): x_request_id = request_headers.get("X-Request-Id", "") x_session_id = request_headers.get("X-Session-Id", "") @@ -756,7 +768,6 @@ async def transfer_to_next_module( self, group_req_objs: Optional[GroupReqObjs] = None, ): - if self.pd_mode.is_P_or_NORMAL(): if not self.args.disable_vision: self.send_to_visual.send_pyobj(group_req_objs.to_group_req_index(), protocol=pickle.HIGHEST_PROTOCOL) @@ -799,7 +810,6 @@ async def _wait_to_token_package( req_status: "ReqStatus", request: Request, ): - event = req_status.event unfinished_count = sampling_params.best_of out_token_counter = 0 @@ -997,6 +1007,7 @@ async def recycle_resource_loop(self): logger.debug(f"httpserver release req_id {req.request_id}, index {req.index_in_shm_mem}") await self.shm_req_manager.async_put_back_req_obj(req) await self.shm_req_manager.async_release_req_index(req.index_in_shm_mem) + self._notify_req_slot_release() await self._release_multimodal_resources(req_status.group_req_objs.multimodal_params) # 先保留这个关键得日志，用于方便定位重构中的问题。 diff --git a/lightllm/server/httpserver/slot_availability.py b/lightllm/server/httpserver/slot_availability.py new file mode 100644 index 0000000000..739166fa6d --- /dev/null +++ b/lightllm/server/httpserver/slot_availability.py @@ -0,0 +1,38 @@ +import asyncio + + +class RequestSlotAvailability: + def __init__(self, enabled=True): + self.enabled = enabled + self.generation = 0 + self.event = asyncio.Event() if enabled else None + + @classmethod + def from_args(cls, args): + return cls( + enabled=( + args.run_mode == "normal" + and args.dp == 1 + and args.nnodes == 1 + and not args.enable_rl + and args.httpserver_workers == 1 + ) + ) + + def notify_release(self): + if self.enabled: + self.generation += 1 + self.event.set() + + async def wait(self, observed_generation, timeout): + if not self.enabled: + await asyncio.sleep(timeout) + return False + if observed_generation != self.generation: + return True + self.event.clear() + try: + await asyncio.wait_for(self.event.wait(), timeout=timeout) + return True + except asyncio.TimeoutError: + return 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..9c8ba806d6 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -172,6 +172,10 @@ def init_model(self, kvargs): with self.model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): manager.finish_state_registration() manager.load_full_att_calibration(args=self.model.args) + if self.args.enable_dp_prompt_cache_fetch: + # 所有 rank 完成模型初始化后再注册共享页池，避免与其他 rank 的加载分配重叠。 + torch.cuda.synchronize() + dist.barrier() manager.initialize_storage( self.model.args, token_capacity=self.model.max_total_token_num, mem_fraction=self.model.mem_fraction ) diff --git a/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py b/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py index b2f3285eee..f2f621b1cf 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py +++ b/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py @@ -104,6 +104,10 @@ def plan(self, decode_reqs: List, origin_batch_size: int) -> SpecDecodePlan: max_batch_size = min(origin_batch_size, available_batch_size) all_reqs_have_proposals = req_num_with_proposals == req_num + if not all_reqs_have_proposals: + self.progress_ema_by_config.clear() + self.prefix_survival_by_depth = [None] * self.max_draft_step + if not self.progress_ema_by_config: # Costs come from CUDA Graph capture, while progress requires one # real verification batch. Keep the available proposal intact until diff --git a/lightllm/server/router/req_queue/__init__.py b/lightllm/server/router/req_queue/__init__.py index d34a0b2933..94ff216527 100644 --- a/lightllm/server/router/req_queue/__init__.py +++ b/lightllm/server/router/req_queue/__init__.py @@ -13,6 +13,14 @@ def _get_req_queue_class(args, router, dp_size_in_node: int): if args.diverse_mode: return ChunkedBeamContinuesBatchQueue + from .chunked_prefill.cached_prefix_queue import ( + CachedPrefixChunkedPrefillQueue, + supports_cached_prefix_reservations, + ) + from lightllm.utils.config_utils import get_model_type + + if get_model_type(args.model_dir) in {"qwen4_exp", "qwen4_exp_text"} and supports_cached_prefix_reservations(args): + return CachedPrefixChunkedPrefillQueue if args.disable_chunked_prefill: # 虽然也使用chuncked prefill queue 但是由于 args.chunked_prefill_size = args.max_req_total_len # 所以调度的实际行为类似过去的 continues batch 调度，所以将两种调度的实现统一为一种实现，减少代码重复。 diff --git a/lightllm/server/router/req_queue/cached_prefix.py b/lightllm/server/router/req_queue/cached_prefix.py new file mode 100644 index 0000000000..8cd7aa5ca5 --- /dev/null +++ b/lightllm/server/router/req_queue/cached_prefix.py @@ -0,0 +1,71 @@ +import os + +from .peak_reservation import Reservation, estimate_peak + + +class CachedPrefixReservations: + def __init__(self, args, dp_index): + self.page_tokens = int(args.linear_att_hash_page_size) + name = os.environ.get("LIGHTLLM_UNIQUE_SERVICE_NAME_ID") + self.domain = (name + ":" + str(dp_index) + ":" + str(self.page_tokens)).encode() if name else None + self.memo, self.by_tuple = {}, {} + self.shared_prefix_tokens = 0 + + def reset(self, requests): + active = {int(request.request_id) for request in requests} + self.memo = {identity: value for identity, value in self.memo.items() if identity in active} + self.by_tuple.clear() + + def bind(self, reservation, request): + self.by_tuple[id(reservation)] = self.prefix(request) + return reservation + + def prefix(self, request): + try: + mm = getattr(request, "multimodal_params", None) + if ( + self.domain is None + or mm is None + or mm.images + or mm.audios + or request.is_paused + or request.is_aborted + or request.sample_params.disable_prompt_cache + ): + return () + length = min(int(request.prompt_cache_len), int(request.input_len), int(request.shm_cur_kv_len)) + hashes = request.linear_att_token_hash_list.get_all() if length >= self.page_tokens else [] + count = min(length // self.page_tokens, len(hashes), 32) + if count <= 0: + return () + hashes = tuple(int(value) & ((1 << 64) - 1) for value in hashes[:count]) + identity = (int(request.request_id), count, hashes, int(request.input_len)) + saved = self.memo.get(int(request.request_id)) + if saved is not None and saved[0] == identity: + return saved[1] + if not hasattr(request, "shm_prompt_ids"): + request.link_prompt_ids_shm_array() + ids = request.shm_prompt_ids.arr + if ids is None or ids.ndim != 1 or ids.dtype.kind not in "iu" or len(ids) < count * self.page_tokens: + return () + pages = tuple( + hashes[index].to_bytes(8, "little") + + ids[index * self.page_tokens : (index + 1) * self.page_tokens].astype(" 0 + ) + + +class CachedPrefixChunkedPrefillQueue(ChunkedPrefillQueue): + def __init__(self, args, router, dp_index, dp_size_in_node): + super().__init__(args, router, dp_index, dp_size_in_node) + self.prefix_reservations = CachedPrefixReservations(args, dp_index) + + def _init_cache_list(self, current_batch, is_busy): + requests = ( + [] + if current_batch is None + else [ + request for request in current_batch.reqs if request.sample_params.suggested_dp_index == self.dp_index + ] + ) + self.prefix_reservations.reset(requests) + super()._init_cache_list(current_batch, is_busy) + + def _get_req_reservation(self, request, is_busy): + reservation = super()._get_req_reservation(request, is_busy) + return self.prefix_reservations.bind(reservation, request) + + def _estimate_peak_tokens(self): + return self.prefix_reservations.calculate(self.cache_len_list) diff --git a/lightllm/server/router/req_queue/chunked_prefill/impl.py b/lightllm/server/router/req_queue/chunked_prefill/impl.py index f69ab11c5e..ab8bebc95d 100644 --- a/lightllm/server/router/req_queue/chunked_prefill/impl.py +++ b/lightllm/server/router/req_queue/chunked_prefill/impl.py @@ -8,7 +8,7 @@ class ChunkedPrefillQueue(BaseQueue): def _init_cache_list(self, current_batch: Batch, is_busy): if current_batch is not None: self.cache_len_list = [ - req.get_tuple_tokens(is_busy, self.router.router_statics.ema_req_out_len) + self._get_req_reservation(req, is_busy) for req in current_batch.reqs if req.sample_params.suggested_dp_index == self.dp_index ] @@ -18,17 +18,10 @@ def _init_cache_list(self, current_batch: Batch, is_busy): # @calculate_time(show=True, min_cost_ms=0.1) def _can_add_new_req(self, req: Req, is_busy): - self.cache_len_list.append( - req.get_tuple_tokens(is_busy, self.router.router_statics.ema_req_out_len) - ) # hard to analysis + self.cache_len_list.append(self._get_req_reservation(req, is_busy)) # hard to analysis self.cache_len_list.sort(key=lambda x: -x[1]) - left_out_len_array = np.array([e[1] for e in self.cache_len_list]) - has_run_len_array = np.array([e[0] for e in self.cache_len_list]) - cum_run_len_array = np.cumsum(has_run_len_array) - size_array = np.arange(1, len(self.cache_len_list) + 1, 1) - - need_max_token_num = (left_out_len_array * size_array + cum_run_len_array).max() + need_max_token_num = self._estimate_peak_tokens() ok_token_num = need_max_token_num < self.max_total_tokens ok_req_num = len(self.cache_len_list) <= self.running_max_req_size @@ -79,14 +72,18 @@ def generate_new_batch(self, current_batch: Batch): def _calcu_batch_token_load_batch_not_none(self, current_batch: Batch): is_busy = self.is_busy() self._init_cache_list(current_batch, is_busy) - if len(self.cache_len_list) != 0: - self.cache_len_list.sort(key=lambda x: -x[1]) - left_out_len_array = np.array([e[1] for e in self.cache_len_list]) - has_run_len_array = np.array([e[0] for e in self.cache_len_list]) - cum_run_len_array = np.cumsum(has_run_len_array) - size_array = np.arange(1, len(self.cache_len_list) + 1, 1) - need_max_token_num = (left_out_len_array * size_array + cum_run_len_array).max() - else: - need_max_token_num = 0 + need_max_token_num = self._estimate_peak_tokens() return (need_max_token_num, need_max_token_num / self.max_total_tokens) + + def _get_req_reservation(self, req, is_busy): + return req.get_tuple_tokens(is_busy, self.router.router_statics.ema_req_out_len) + + def _estimate_peak_tokens(self): + if not self.cache_len_list: + return 0 + self.cache_len_list.sort(key=lambda entry: -entry[1]) + remaining = np.array([entry[1] for entry in self.cache_len_list]) + resident = np.array([entry[0] for entry in self.cache_len_list]) + count = np.arange(1, len(self.cache_len_list) + 1) + return (remaining * count + np.cumsum(resident)).max() diff --git a/lightllm/server/router/req_queue/peak_reservation.py b/lightllm/server/router/req_queue/peak_reservation.py new file mode 100644 index 0000000000..9df7ab6144 --- /dev/null +++ b/lightllm/server/router/req_queue/peak_reservation.py @@ -0,0 +1,52 @@ +from dataclasses import dataclass +from typing import Optional, Sequence, Tuple + + +@dataclass(frozen=True) +class Reservation: + base_tokens: int + remaining_tokens: int + prefix_pages: Tuple[bytes, ...] = () + cache_domain: Optional[bytes] = None + + +@dataclass(frozen=True) +class Peak: + legacy_tokens: int + shared_tokens: int + max_prefix_credit: int + + +def estimate_peak(entries: Sequence[Reservation], page_tokens: int) -> Peak: + if type(page_tokens) is not int or page_tokens <= 0: + raise ValueError("Page token count must be a positive integer") + for entry in entries: + if ( + type(entry.base_tokens) is not int + or entry.base_tokens < 0 + or type(entry.remaining_tokens) is not int + or entry.remaining_tokens < 0 + ): + raise ValueError("Reservation counts must be nonnegative integers") + if not isinstance(entry.prefix_pages, tuple) or any(type(page) is not bytes for page in entry.prefix_pages): + raise ValueError("Prefix pages must have exact immutable identities") + if entry.cache_domain is not None and type(entry.cache_domain) is not bytes: + raise ValueError("An explicit shared domain must have immutable identity") + if len(entry.prefix_pages) * page_tokens > entry.base_tokens: + raise ValueError("Known cached pages exceed the resident base") + seen = set() + base = credit = legacy = shared = 0 + for count, entry in enumerate(sorted(entries, key=lambda entry: -entry.remaining_tokens), 1): + base += entry.base_tokens + if entry.cache_domain is not None: + prefix = (entry.cache_domain,) + for page in entry.prefix_pages: + prefix += (page,) + if prefix in seen: + credit += page_tokens + else: + seen.add(prefix) + projected = base + count * entry.remaining_tokens + legacy = max(legacy, projected) + shared = max(shared, projected - credit) + return Peak(legacy, shared, credit) diff --git a/lightllm/server/tokenizer.py b/lightllm/server/tokenizer.py index 12ddbb1c68..1fb836998d 100644 --- a/lightllm/server/tokenizer.py +++ b/lightllm/server/tokenizer.py @@ -124,6 +124,15 @@ def get_tokenizer( tokenizer = QWen3VLTokenizer( tokenizer=tokenizer, image_processor=processor.image_processor, model_cfg=model_cfg ) + elif model_type == "qwen4_exp" and "vision_config" in model_cfg: + from ..models.qwen4_exp.image_processing import Qwen4ExpImageProcessor + from ..models.qwen4_exp.tokenizer import Qwen4ExpTokenizer + + tokenizer = Qwen4ExpTokenizer( + tokenizer=tokenizer, + image_processor=Qwen4ExpImageProcessor.from_pretrained(tokenizer_name), + model_cfg=model_cfg, + ) elif model_type in ["qwen3_5", "sensenova_loop", "qwen3_5_moe"] and "vision_config" in model_cfg: from ..models.qwen3_5.model import QWen3_5Tokenizer from transformers import AutoImageProcessor diff --git a/lightllm/server/visualserver/model_infer/model_rpc.py b/lightllm/server/visualserver/model_infer/model_rpc.py index 6451e70008..b8751db58b 100644 --- a/lightllm/server/visualserver/model_infer/model_rpc.py +++ b/lightllm/server/visualserver/model_infer/model_rpc.py @@ -92,6 +92,10 @@ def exposed_init_model(self, kvargs): self.model = ( Qwen2_5_VisionTransformerPretrainedModel(kvargs, **model_cfg["vision_config"]).eval().bfloat16() ) + elif self.model_type == "qwen4_exp": + from lightllm.models.qwen4_exp.vision import Qwen4ExpVisionTransformer + + self.model = Qwen4ExpVisionTransformer(kvargs, **model_cfg["vision_config"]).eval().bfloat16() elif self.model_type in ["qwen3_vl", "qwen3_vl_moe", "qwen3_5", "sensenova_loop", "qwen3_5_moe"]: self.model = ( Qwen3VisionTransformerPretrainedModel(kvargs, **model_cfg["vision_config"]).eval().bfloat16() diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 1b0ce53879..187af4706d 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -216,6 +216,8 @@ def auto_set_fused_shared_experts(args) -> None: "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text", + "qwen4_exp", + "qwen4_exp_text", } if model_type not in supported_model_types: logger.info(f"do not enable fused shared experts: unsupported model_type={model_type}") @@ -288,7 +290,15 @@ def get_eos_token_ids(model_path: str) -> Optional[List[int]]: try: config_json = get_config_json(model_path) model_type = config_json.get("model_type") or config_json.get("text_config", {}).get("model_type") - if model_type in {"qwen3_5", "sensenova_loop", "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text"}: + if model_type in { + "qwen3_5", + "sensenova_loop", + "qwen3_5_text", + "qwen3_5_moe", + "qwen3_5_moe_text", + "qwen4_exp", + "qwen4_exp_text", + }: from transformers import AutoTokenizer eos_token_ids = [] @@ -419,6 +429,8 @@ def has_vision_module(model_path: str) -> bool: return True elif model_type == "neo_chat": return True + elif model_type == "qwen4_exp": + return model_cfg.get("vision_config") is not None elif model_type in ["qwen3_5", "sensenova_loop", "qwen3_5_moe"]: return True else: @@ -464,6 +476,8 @@ def is_linear_att_mixed_model(model_path: str) -> bool: "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text", + "qwen4_exp", + "qwen4_exp_text", "glm5_next", "glm5_next_text", ] @@ -559,6 +573,8 @@ def get_reasoning_parser_for_model(model_path: str) -> Optional[str]: "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text", + "qwen4_exp", + "qwen4_exp_text", ]: return "qwen3" diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py new file mode 100644 index 0000000000..f4928a9017 --- /dev/null +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -0,0 +1,162 @@ +"""Measure a full KDA state cycle: all layers, acceptance, folds and export. + +This synthetic kernel benchmark excludes convolutions, model projections and +HTTP scheduling. Use a matched service workload for end-to-end conclusions. +""" + +import argparse +import gc +import json +import statistics +from pathlib import Path + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda + + +def ReplaySSMExecutor(*args, **kwargs): + kda = kwargs.pop("kda", False) + if kda: + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner + else: + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner + + return _history_owner(*args, **kwargs) + + +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs, replay_layer + + +def workload(args, batch, width, dtype, mode, acceptance): + layers, heads, dim = args.layers, args.heads, 128 + shape = (1, batch * width, heads, dim) if width > 1 else (batch, 1, heads, dim) + packed = torch.randn(batch * width, heads * dim * 3, dtype=torch.bfloat16, device="cuda") + q, k, v = [x.view(shape) for x in packed.split(heads * dim, -1)] + gate = torch.randn(batch * width, heads * dim, dtype=q.dtype, device="cuda") - 3 + beta = torch.randn(batch * width, heads, dtype=q.dtype, device="cuda") + alog = torch.zeros(heads, device="cuda") + bias = torch.zeros(heads * dim, device="cuda") + reqs = torch.arange(batch, dtype=torch.int32, device="cuda") + cu = torch.arange(batch + 1, dtype=torch.int32, device="cuda") * width if width > 1 else None + counts = reqs % width + 1 if acceptance == "partial" else torch.full_like(reqs, width) + accepted = torch.cat((counts - 1, counts.new_zeros(1))) + snapshot = torch.empty(layers, heads, dim, dim, dtype=dtype, device="cuda") + slots = (batch + 1) * (width if mode == "native" else 1) + state = torch.randn(layers, slots, heads, dim, dim, dtype=dtype, device="cuda") * 0.01 + ssm_backend = None + if mode != "native": + capacity = max(8, triton.next_power_of_2(width)) + ssm_backend = ReplaySSMExecutor(state, capacity, width, q.dtype, kda=True, compute_mode=mode) + if ssm_backend is not None: + ssm_backend._config_is_fixed = True + indexes = reqs[:, None] * width + torch.arange(width, device="cuda", dtype=torch.int32) + if width == 1: + indexes = reqs + bytes_used = ( + sum(x.numel() * x.element_size() for x in vars(ssm_backend).values() if isinstance(x, torch.Tensor)) + if ssm_backend + else state.numel() * state.element_size() + ) + + native_index = 0 if acceptance == "partial" else width - 1 + + def cycle(): + for _ in range(args.rounds): + if ssm_backend is not None: + ssm_backend.prepare_decode(reqs, cu) + for layer in range(layers): + if ssm_backend is None: + fused_recurrent_kda( + q, + k, + v, + gate.view(*shape[:2], -1), + beta.view(*shape[:2], heads), + alog, + bias, + state[layer], + indexes, + cu_seqlens=cu, + num_accepted_tokens=counts if cu is not None else None, + ) + else: + replay_layer(ssm_backend, layer, q, k, v, gate, beta, alog, bias, reqs, cu) + if ssm_backend is not None: + ssm_backend.accept_updates(reqs, accepted) + if isinstance(ssm_backend, ReplaySSMExecutor): + ssm_backend._materialize_accepted_state(reqs[:1], snapshot, snapshot=True) + ssm_backend.merge_accepted_updates(reqs) + else: + snapshot.copy_(state[:, native_index] if ssm_backend is None else state[:, 0]) + + return ssm_backend, cycle, bytes_used + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--batch", type=int, nargs="+", default=[16, 64, 192]) + p.add_argument("--width", type=int, nargs="+", default=[1, 3]) + p.add_argument("--dtype", nargs="+", choices=["float32", "bfloat16"], default=["float32", "bfloat16"]) + p.add_argument("--layers", type=int, default=34) + p.add_argument("--heads", type=int, default=8) + p.add_argument("--rounds", type=int, default=8) + p.add_argument("--repeats", type=int, default=5) + p.add_argument("--tune", action="store_true") + p.add_argument("--output", type=Path, required=True) + args = p.parse_args() + torch.manual_seed(53) + rows = [] + report = dict( + settings={**vars(args), "output": str(args.output)}, + device=torch.cuda.get_device_name(), + torch=torch.__version__, + triton=triton.__version__, + results=rows, + ) + for dtype_name in args.dtype: + for batch in args.batch: + for width in args.width: + for acceptance in ["full", "partial"] if width > 1 else ["full"]: + for mode in ["native", "inline", "precompute"] if width > 1 else ["native", "inline"]: + ssm_backend, cycle, size = workload( + args, batch, width, getattr(torch, dtype_name), mode, acceptance + ) + candidates = [] + if ssm_backend is not None: + for config in get_configs() if args.tune else [{"BV": 32, "num_warps": 2, "num_stages": 1}]: + if mode == "precompute" and config["BV"] < 16: + continue + ssm_backend.run_config = config + cycle() + ms = triton.testing.do_bench_cudagraph(cycle, rep=50) + candidates.append(dict(config=config, ms=ms)) + ssm_backend.run_config = min(candidates, key=lambda r: r["ms"])["config"] + cycle() + samples = [ + triton.testing.do_bench_cudagraph(cycle, rep=100) / args.rounds for _ in range(args.repeats) + ] + row = dict( + batch=batch, + width=width, + dtype=dtype_name, + acceptance=acceptance, + mode=mode, + cache_len=getattr(ssm_backend, "capacity", None), + state_history_bytes=size, + ms_per_round=statistics.median(samples), + samples_ms=samples, + config=dict(ssm_backend.run_config) if ssm_backend and ssm_backend.run_config else None, + candidates=candidates, + ) + rows.append(row) + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(row), flush=True) + del ssm_backend, cycle + gc.collect() + torch.cuda.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/kernels/benchmark_replayssm_service.py b/test/benchmark/kernels/benchmark_replayssm_service.py new file mode 100644 index 0000000000..a1cee4826a --- /dev/null +++ b/test/benchmark/kernels/benchmark_replayssm_service.py @@ -0,0 +1,119 @@ +"""Compare streamed requests on baseline and optimized services with identical warmup. + +Reports request timings and token IDs for an exact cross-revision comparison. +""" + +import argparse +import concurrent.futures +import json +import statistics +import time +import threading +from pathlib import Path + +import requests +from transformers import AutoTokenizer + +p = argparse.ArgumentParser() +p.add_argument("--output", required=True) +p.add_argument("--batch", type=int, nargs="+", default=[1, 8, 16]) +p.add_argument("--repeats", type=int, default=3) +p.add_argument("--tokens", type=int, default=64) +p.add_argument("--warmups", type=int, default=1) +p.add_argument("--model-dir", required=True) +p.add_argument("--url", default="http://127.0.0.1:18785/generate_stream") +args = p.parse_args() +url = args.url +tok = AutoTokenizer.from_pretrained(args.model_dir, trust_remote_code=True) +prefix = tok.encode( + "Describe how a computer executes a program, from source code to machine instructions. " + "Explain the role of memory, processors, and operating systems. " * 100 +)[:512] + + +def prepare_prompt(index, rep): + suffix = tok.encode( + f"Example number {rep * 1000 + index}. Discuss this scenario in detail. " + + ("An engineer measures performance and validates numerical correctness. " * 100) + )[:512] + return prefix + suffix + + +def stream_request(index, rep, ids, barrier, cohort_start): + payload = { + "inputs": ids, + "parameters": { + "do_sample": False, + "disable_prompt_cache": True, + "ignore_eos": True, + "max_new_tokens": args.tokens, + }, + } + barrier.wait() + start = time.perf_counter() + times = [] + tokens = [] + texts = [] + finished = False + with requests.post(url, json=payload, stream=True, timeout=300) as response: + response.raise_for_status() + # Consume each HTTP chunk; byte-sized iteration can saturate the client CPU. + for line in response.iter_lines(chunk_size=None): + if not line.startswith(b"data:"): + continue + item = json.loads(line[5:]) + assert "error" not in item, item + token = item.get("token", {}) + if token.get("id") is not None: + times.append(time.perf_counter() - cohort_start) + tokens.append(token["id"]) + texts.append(token.get("text", "")) + finished = item.get("finished", finished) + assert finished and len(tokens) == args.tokens, (finished, len(tokens)) + return dict( + index=index, + rep=rep, + prompt_len=len(ids), + token_ids=tokens, + text="".join(texts), + token_times_s=times, + ttft_ms=(times[0] - (start - cohort_start)) * 1000, + tpot_ms=(times[-1] - times[0]) * 1000 / max(1, len(times) - 1), + elapsed_s=time.perf_counter() - start, + ) + + +def run_concurrent_requests(batch, rep): + prompts = [prepare_prompt(i, rep) for i in range(batch)] + barrier = threading.Barrier(batch) + start = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=batch) as pool: + samples = list(pool.map(lambda i: stream_request(i, rep, prompts[i], barrier, start), range(batch))) + elapsed = time.perf_counter() - start + begin = max(r["token_times_s"][0] for r in samples) + 0.5 + end = min(r["token_times_s"][-1] for r in samples) - 0.5 + count = sum(sum(begin <= t <= end for t in r["token_times_s"]) for r in samples) + return dict( + batch=batch, + rep=rep, + elapsed_s=elapsed, + tokens_per_s=sum(len(r["token_ids"]) for r in samples) / elapsed, + steady_tokens_per_s=count / (end - begin) if end > begin else None, + steady_window_s=max(0.0, end - begin), + ttft_ms=statistics.median(r["ttft_ms"] for r in samples), + tpot_ms=statistics.median(r["tpot_ms"] for r in samples), + requests=samples, + ) + + +# Warm every measured graph bucket before timing, with the same requests on both revisions. +for batch in args.batch: + for warm in range(args.warmups): + run_concurrent_requests(batch, 99 + warm) +rows = [] +for batch in args.batch: + for rep in range(args.repeats): + row = run_concurrent_requests(batch, rep) + rows.append(row) + Path(args.output).write_text(json.dumps(rows, indent=2)) + print(json.dumps({k: v for k, v in row.items() if k != "requests"}), flush=True) diff --git a/test/benchmark/qwen38/README.agentx.md b/test/benchmark/qwen38/README.agentx.md new file mode 100644 index 0000000000..ff05d234a0 --- /dev/null +++ b/test/benchmark/qwen38/README.agentx.md @@ -0,0 +1,171 @@ +# Repeatable AgentX control replay + +Use this protocol before comparing or tuning engines. A run that completes is not +automatically a valid performance result: all request audits and the three-trial +repeatability gate must pass. + +## What is measured + +The same pinned AgentX installation drives every engine. `prepare_agentx.py` uses +its Weka loader and synthetic prompt generator offline, then freezes full request +bodies and synthetic conversation history into `inputs.json`. Measurement uses +the installation's AIPerf `inputs_json` loader, which sends these payloads verbatim. +Live responses cannot change subsequent inputs. + +The default workload is 32 independent streams, eight sequential turns per stream, +256 requests in total, with concurrency 32. Output budgets come from the source +trace, with `ignore_eos=true`, temperature zero and seed 42. Each stream runs once +and drains completely. There is no duration cutoff, resampling, request clipping, +think delay or DAG join. Prefix reuse within and across the frozen streams remains +enabled. Tail drain is included in throughput, so concurrency falls as streams +finish; this is a fixed-work control, not a sustained-capacity measurement. + +This is **Weka-derived control replay**, not the official AgentX score. The +official scenario retains trace delays, branching and response-dependent history; +even the same seed does not make those future requests byte-identical. Neither +mode establishes real business quality or production capacity. Do not mix their +results in one ranking. + +## Frozen inputs + +- Tool: `agentx-harness==1.0.13`. The manifest records all installed dependency + versions and hashes every Python source file in `aiperf` and `agentx_cli.py`. + Use Python 3.11--3.13 and a dedicated environment. Install + `agentx-harness==1.0.13` before preparation; thereafter reuse the same environment + or restore the exact dependency versions in the manifest. A matching version + number alone is insufficient: the runner also requires the source hashes and + complete installed distribution list to match. +- Dataset: `semianalysisai/cc-traces-weka-062126-256k`, revision + `8fecd2fc56694469f758f0afbbb6335ad3043740`, `traces.jsonl` SHA-256 + `e39cd2ff3eba21d4a3664be51da743ac3d2149a1933898cafc7bfeac8147eeef`. +- Default model name: `Qwen3.8-Flash-Next`. The manifest hashes configuration, + tokenizer, chat template, acquisition metadata and weight index, and records + every weight shard's size. Weight contents are not rehashed by this tool; use + the same verified, read-only checkpoint mount for all engines. +- Every selected request is checked with the checkpoint's chat template against + the 262,144-token prompt-plus-output limit. An overflow fails preparation. +- `sources.json` maps each turn to its original Weka trace and request index. + Frozen histories are synthetic reconstructions, not real customer text. + +Prepare once and reuse that exact output directory for every engine. Never +regenerate a different workload for a slower engine. Preparation is deterministic; +verify a second independent generation has identical `inputs.json` bytes. + +```sh +exp -m "Freeze AgentX control; record source checkpoint and purpose" \ + /path/to/agentx-venv/bin/python -m test.benchmark.qwen38.prepare_agentx \ + --source /path/to/snapshots/8fecd2fc56694469f758f0afbbb6335ad3043740/traces.jsonl \ + --model-directory /path/to/verified/model \ + --output /path/to/new/frozen-workload +``` + +## Server and trial procedure + +Review previous `/data6/exp_logs` records (on this H100 host, +`/home/devsft/experiments/runs`) before each experiment. Launch the engine through +`exp` and record its source revision, image digest, complete command, checkpoint, +GPU assignment, precision, KV format, EP/TP, MTP and cache settings. Engine tuning +must use a separate screening set; freeze the final configuration before a suite. + +The runner is restricted to this task's isolated eight-GPU container, labelled +`qwen38.task=a3844d54-8a2b-4532-bf54-f657c9ca72a0`, and the localhost endpoint +`http://127.0.0.1:28081`. It refuses another running Docker container or a foreign +GPU compute process at phase boundaries. Use a dedicated host; stop image pulls, +image imports, builds, weight copies and other benchmark clients before starting. +Do not change clocks, power limits or engine configuration between trials. + +Enable and verify the engine's cache-reset API before measurement: + +| Engine | Required API | Required confirmation | Runtime status | +| --- | --- | --- | --- | +| vLLM 0.33 | `POST /reset_prefix_cache` | `{"success":true}` | Validated on H100 | +| LightLLM | `POST /flush_cache`, requires `--enable_rl` | `{"success":true}` | Adapter awaits runtime validation | +| SGLang 0.5.23 | `POST /flush_cache` | HTTP 200, `Cache flushed.` text or empty/null/true body | Streaming/reset adapter validated on H100 | + +The tested vLLM image requires `VLLM_SERVER_DEV_MODE=1` to expose the reset +endpoint. Bind this benchmark service to localhost. A missing endpoint, unexpected +response or failed reset aborts the suite; it is never silently skipped. + +Each of at least three trials executes: + +1. Confirm the same container and frozen tool/checkpoint/workload. +2. Reset the server cache and require success. +3. Take three idle GPU samples, ten seconds apart: utilization at most 5%, + temperature at most 60 C, all eight GPUs present, no foreign compute PID. + Wait up to five minutes for three consecutive qualifying samples, retaining + every check. Missing GPUs or a foreign compute PID abort immediately. +4. Replay and audit the entire fixed workload as warmup. +5. Reset cache again, repeat the idle checks, then replay and audit the same full + workload for measurement. + +Warmup is identical in every trial; no implicit first-run-only pressure warmup is +used. Only the measurement phase enters the reported three-trial statistics. + +```sh +exp -m "Qualify repeatability; record engine/image/flags, hardware, workload and expected question" \ + /path/to/agentx-venv/bin/python -m test.benchmark.qwen38.run_agentx \ + --engine vllm \ + --container owned-engine-container \ + --workload /path/to/frozen-workload \ + --model-directory /path/to/verified/model \ + --output /path/to/new/suite --repeats 3 +``` + +For SGLang or LightLLM, change `--engine` and the container only after independently +validating the reset and OpenAI streaming adapters. Keep the same model name, +tokenizer, full request file, load and measurement rules. Disclose precision and +speculation differences rather than treating them as equivalent accuracy. + +## Acceptance and evidence + +`audit_agentx.py` rejects missing/duplicate/changed payloads, incorrect conversation +turn mappings or order, overlapping turns in one stream, HTTP errors, tool errors, +cancellations, skipped contexts, missing SSE `[DONE]` or usage, non-length finish, +and output counts different from the frozen budget. It independently computes +output tokens / (last response end - first request start), requiring agreement +with the tool's throughput within 2%. Per-request prompt counts must also match +across repeated trials. + +The numerator is server `usage.completion_tokens`, including reasoning tokens. +AgentX `output_token_count` counts visible output separately and must not be used +as the total for reasoning models. Reconcile `usage_completion_tokens` with raw +SSE usage and use the independently counted throughput for engine comparisons. +SGLang 0.5.23 can report reasoning tokens above completion tokens; AgentX 1.0.13 +then slightly inflates `output_sequence_length` and its derived throughput. Keep +the tool and server unchanged during a suite, disclose this discrepancy and +retain the 2% independent cross-check. + +All trials must pass these predeclared limits: + +| Metric | Sample standard deviation / mean (CV) | (max - min) / mean | +| --- | --- | --- | +| Output token throughput | <= 5% | <= 10% | +| Request throughput | <= 5% | <= 10% | +| TTFT p50 | <= 10% | <= 20% | + +The limits are protocol choices, not guarantees provided by AgentX. Do not drop +the slowest trial, select the best trial or loosen limits after viewing results. +Failure exits nonzero and preserves evidence. Investigate it before ranking +engines. Passing establishes repeatability for this workload and configuration; +it does not establish repeatability across service restarts or hosts. Qualify a +fresh-launch suite separately before publishing broader claims. + +`suite.json` includes the manifest, protocol hashes, container image/command, +reset replies, idle samples, per-phase audits and final variance gate. Each trial +retains its config, exact command, PID, log, raw request/SSE exports, tool summary +and server metrics before/after. The suite also keeps `nvidia-smi -q`, five-second +GPU utilization/memory/power/temperature samples and server logs. Review telemetry +and host activity for interference; phase-boundary checks cannot prove that no +short-lived external process ran during measurement. + +Recompute the variance gate without model requests: + +```sh +python -m test.benchmark.qwen38.audit_agentx \ + --suite /path/to/suite/suite.json --workload /path/to/frozen-workload +``` + +For engine comparisons, report all trial values, mean, CV, range, lengths, errors, +TTFT p50/p95 and evidence locations. Alternate engine order in a later independent +validation block to check order effects. Older runs with different warmup, +cache or history conditions must remain excluded. diff --git a/test/benchmark/qwen38/audit_agentx.py b/test/benchmark/qwen38/audit_agentx.py new file mode 100644 index 0000000000..eb60fbd46d --- /dev/null +++ b/test/benchmark/qwen38/audit_agentx.py @@ -0,0 +1,158 @@ +import argparse +from collections import Counter +import json +import math +from pathlib import Path +import statistics + +from test.benchmark.qwen38.prepare_agentx import canonical, digest + + +def audit(directory, inputs): + expected = Counter(digest(canonical(p)) for s in inputs["data"] for p in s["payloads"]) + turns = {(s["session_id"], i): p for s in inputs["data"] for i, p in enumerate(s["payloads"])} + observed_turns, prompt_counts, intervals = set(), {}, {} + seen, prompt_tokens, output_tokens, starts, ends = Counter(), 0, 0, [], [] + merged = directory / "profile_export_raw.jsonl" + files = [merged] if merged.is_file() else sorted((directory / "raw_records").glob("*.jsonl")) + if not files: + raise ValueError("Missing raw request/SSE export shards") + for row in raw_rows(files): + metadata = row["metadata"] + if metadata["benchmark_phase"] != "profiling": + raise ValueError("Unexpected warmup records in a standalone measurement") + if ( + row["status"] != 200 + or row.get("error") + or metadata.get("was_cancelled") + or metadata.get("context_overflow_skip") + ): + raise ValueError("Failed, cancelled or skipped request") + payload = row["payload"] + key = (metadata["conversation_id"], metadata["turn_index"]) + if key not in turns or key in observed_turns or canonical(payload) != canonical(turns[key]): + raise ValueError("Missing, duplicate or changed session turn") + observed_turns.add(key) + seen[digest(canonical(payload))] += 1 + usage, done, reason = None, 0, None + for response in row["responses"]: + for packet in response.get("packets", []): + if packet.get("name") != "data": + continue + if done: + raise ValueError("SSE data after DONE") + if packet["value"] == "[DONE]": + done += 1 + continue + body = json.loads(packet["value"]) + usage = body.get("usage") or usage + for choice in body.get("choices", []): + reason = choice.get("finish_reason") or reason + if done != 1 or reason != "length" or not usage: + raise ValueError("Missing usage, DONE or length termination") + if type(usage.get("completion_tokens")) is not int or usage["completion_tokens"] != payload["max_tokens"]: + raise ValueError("Server output count differs from fixed work") + if type(usage.get("prompt_tokens")) is not int or usage["prompt_tokens"] <= 0: + raise ValueError("Missing actual prompt count") + prompt_tokens += usage["prompt_tokens"] + output_tokens += usage["completion_tokens"] + prompt_counts[key] = usage["prompt_tokens"] + if metadata["request_start_ns"] >= metadata["request_end_ns"]: + raise ValueError("Invalid request interval") + intervals[key] = (metadata["request_start_ns"], metadata["request_end_ns"]) + starts.append(metadata["request_start_ns"]) + ends.append(metadata["request_end_ns"]) + if seen != expected: + raise ValueError("Missing, duplicate or changed payloads") + for (session, turn), (start, end) in intervals.items(): + if turn and start < intervals[(session, turn - 1)][1]: + raise ValueError("Session turns overlapped or ran out of order") + summary = json.loads((directory / "profile_export_aiperf.json").read_text()) + if summary.get("error_summary") or summary.get("was_cancelled"): + raise ValueError("AgentX reported errors or cancellation") + if summary["request_count"]["avg"] != sum(expected.values()): + raise ValueError("AgentX request count mismatch") + if summary["usage_completion_tokens"]["sum"] != output_tokens: + raise ValueError("AgentX output token total differs from SSE usage") + if summary["usage_prompt_tokens"]["sum"] != prompt_tokens: + raise ValueError("AgentX prompt token total differs from SSE usage") + seconds = (max(ends) - min(starts)) / 1e9 + if not math.isfinite(seconds) or seconds <= 0: + raise ValueError("Invalid request wall time") + throughput = output_tokens / seconds + reported = summary["output_token_throughput"]["avg"] + if not math.isfinite(reported) or abs(reported / throughput - 1) > 0.02: + raise ValueError("AgentX throughput differs from independent count / request wall time by >2%") + request_tps = summary["request_throughput"]["avg"] + if not math.isfinite(request_tps) or abs(request_tps / (len(observed_turns) / seconds) - 1) > 0.02: + raise ValueError("AgentX request throughput differs from independent count / request wall time by >2%") + return dict( + payloads_sha256=digest(canonical(sorted(expected.items()))), + prompt_counts_sha256=digest(canonical(sorted(prompt_counts.items()))), + requests=sum(expected.values()), + prompt_tokens=prompt_tokens, + output_tokens=output_tokens, + request_wall_s=seconds, + output_tps=reported, + independent_output_tps=throughput, + request_tps=request_tps, + ttft_p50_ms=summary["time_to_first_token"]["p50"], + ttft_p95_ms=summary["time_to_first_token"]["p95"], + ) + + +def raw_rows(files): + for path in files: + with path.open() as handle: + for line in handle: + yield json.loads(line) + + +def repeatability(runs): + if len(runs) < 3: + raise ValueError("Require at least three audited runs") + for key in ("payloads_sha256", "prompt_counts_sha256", "requests", "prompt_tokens", "output_tokens"): + if len({r[key] for r in runs}) != 1: + raise ValueError(f"Run workload mismatch: {key}") + metrics = {} + for key, limit in (("output_tps", 0.05), ("request_tps", 0.05), ("ttft_p50_ms", 0.10)): + values = [r[key] for r in runs] + if any(not math.isfinite(v) or v <= 0 for v in values): + raise ValueError(f"Invalid metric: {key}") + average = statistics.mean(values) + metrics[key] = dict( + values=values, + mean=average, + cv=statistics.stdev(values) / average, + relative_range=(max(values) - min(values)) / average, + cv_limit=limit, + ) + passed = all(m["cv"] <= m["cv_limit"] and m["relative_range"] <= 2 * m["cv_limit"] for m in metrics.values()) + return dict(status="passed" if passed else "nonrepeatable", runs=len(runs), metrics=metrics) + + +def main(): + parser = argparse.ArgumentParser(description="Reaudit fixed AgentX exports and trial variance") + parser.add_argument("--suite", type=Path, required=True) + parser.add_argument("--workload", type=Path, required=True) + args = parser.parse_args() + suite = json.loads(args.suite.read_text()) + raw = (args.workload / "inputs.json").read_bytes() + if digest(raw) != suite["manifest"]["inputs_sha256"]: + raise ValueError("Workload differs from recorded suite") + inputs = json.loads(raw) + runs = [] + for trial in suite["runs"]: + for phase in ("warmup", "measure"): + current = audit(args.suite.parent / f"trial-{trial['trial']:02}" / phase, inputs) + if current != trial[phase + "_audit"]: + raise ValueError("Retained raw audit differs from the suite") + if phase == "measure": + runs.append(current) + result = repeatability(runs) + print(json.dumps(result, indent=2)) + raise SystemExit(0 if result["status"] == "passed" else 1) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/benchmark_gated_residual.py b/test/benchmark/qwen38/benchmark_gated_residual.py new file mode 100644 index 0000000000..0e3c4fa039 --- /dev/null +++ b/test/benchmark/qwen38/benchmark_gated_residual.py @@ -0,0 +1,93 @@ +import argparse +import json +import os + +import torch +import torch.nn.functional as F +import triton +from triton.testing import do_bench_cudagraph + +from lightllm.common.gated_residual import GatedResidual + + +def eager(component, residual, update, precise_projections=False): + def project(inputs, weight): + if precise_projections: + return F.linear(inputs.double(), weight.double()).to(inputs.dtype) + return F.linear(inputs, weight) + + grouped = residual.float().unflatten(-1, (component.branches, -1)) + normalized = grouped * torch.rsqrt(grouped.square().mean(-1, keepdim=True) + component.eps) + normalized = (normalized.flatten(-2) * (1 + component.norm_weight.float())).to(residual.dtype) + projected = F.silu(project(normalized, component.down_weight) / component.branches) + mix = torch.sigmoid(project(projected, component.up_weight)) + mixed = (normalized * mix).unflatten(-1, (component.branches, -1)).mean(-2) + gates = 2 * torch.sigmoid(project(normalized, component.inject_weight) / component.branches) + return mixed, residual + (update.unsqueeze(-2) * gates.unsqueeze(-1)).flatten(-2) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokens", type=int, nargs="+", default=[1, 16, 128, 1024]) + parser.add_argument("--repetitions", type=int, default=3) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Wrap this benchmark with exp -m "purpose"') + torch.manual_seed(37) + dim, branches, rank = 2560, 4, 320 + width = dim * branches + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) + + component = GatedResidual( + rand(width) * 0.1, + rand(rank, width) / width ** 0.5, + rand(width, rank) / rank ** 0.5, + rand(branches, width) / width ** 0.5, + branches, + ) + print( + json.dumps( + { + "torch": torch.__version__, + "triton": triton.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "dtype": "bfloat16", + "hidden_size": dim, + "branches": branches, + "rank": rank, + "timing": "CUDA graph, median of repetitions", + "validation": "eager equations with FP64 projections quantized at original dtype boundaries", + "eager_timing": "original input-dtype projection; batch-dependent rounding recorded separately", + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + } + ) + ) + for tokens in args.tokens: + residual, update = rand(tokens, width), rand(tokens, dim) + + def fused(): + mixed, gates = component.read(residual) + return mixed, component.write(residual, update, gates) + + def reference(): + return eager(component, residual, update) + + actual = fused() + expected = eager(component, residual, update, precise_projections=True) + for value, oracle in zip(actual, expected): + torch.testing.assert_close(value, oracle, atol=2e-3, rtol=torch.finfo(residual.dtype).eps) + results = { + "eager_max_abs": [float((value - old).float().abs().max()) for value, old in zip(actual, reference())] + } + for name, func in [("eager", reference), ("fused", fused)]: + samples = [do_bench_cudagraph(func, rep=100) * 1000 for _ in range(args.repetitions)] + results[name + "_us"] = sorted(samples)[len(samples) // 2] + results[name + "_samples_us"] = samples + print(json.dumps({"tokens": tokens, **results, "speedup": results["eager_us"] / results["fused_us"]})) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/benchmark_lm_head.py b/test/benchmark/qwen38/benchmark_lm_head.py new file mode 100644 index 0000000000..422c9018c2 --- /dev/null +++ b/test/benchmark/qwen38/benchmark_lm_head.py @@ -0,0 +1,111 @@ +import argparse +import importlib.metadata +import json +import os +from pathlib import Path +import subprocess + +import torch +import triton +from safetensors import safe_open + +from check_generation import REVISION, verified_checkpoint + + +@torch.inference_mode() +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument( + "--tp", type=int, default=4, help="Vocabulary partition count; this process runs rank zero only" + ) + parser.add_argument("--rows", type=int, nargs="+", default=[1, 3, 5, 7, 8, 9, 24, 64, 128]) + parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16") + parser.add_argument("--samples", type=int, default=3) + parser.add_argument("--rep-ms", type=int, default=40) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + raise RuntimeError("Run through exp") + if min(args.rows + [args.tp, args.samples, args.rep_ms]) < 1: + parser.error("Partition count, rows, samples and repeat duration must be positive") + verified = verified_checkpoint(args.model_directory) + checkpoint = json.loads((args.model_directory / "config.json").read_text()) + config = checkpoint.get("text_config", checkpoint) + hidden, vocab = config["hidden_size"], config["vocab_size"] + if args.tp > vocab: + parser.error("Partition count cannot exceed vocabulary size") + + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils.envs_utils import get_env_start_args + + start = StartArgs() + start.hardware_platform, start.data_type, start.tp, start.dp = "cuda", args.dtype, args.tp, 1 + os.environ["LIGHTLLM_START_ARGS"] = json.dumps(vars(start)) + os.environ.update( + LIGHTLLM_DP_WORLD_SIZE=str(args.tp), LIGHTLLM_CURRENT_RANK_IN_DP="0", LIGHTLLM_CURRENT_DEVICE_ID="0" + ) + get_env_start_args.cache_clear() + from lightllm.common.basemodel.layer_weights.meta_weights import LMHeadWeight + from lightllm.models.qwen4_exp.layer_weights.lm_head_weight import Qwen4ExpLMHeadWeight + + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 + baseline = LMHeadWeight(hidden, vocab, "lm_head.weight", dtype) + index = json.loads((args.model_directory / "model.safetensors.index.json").read_text())["weight_map"] + filename = index["lm_head.weight"] + with safe_open(args.model_directory / filename, framework="pt", device="cpu") as source: + baseline.load_hf_weights({"lm_head.weight": source.get_tensor("lm_head.weight")}) + assert baseline.verify_load() + optimized = Qwen4ExpLMHeadWeight(hidden, vocab, "lm_head.weight", dtype, embedding_weight=baseline) + assert optimized.verify_load() and optimized.weight is baseline.weight + torch.manual_seed(3107) + inputs = torch.randn(max(args.rows), hidden, device="cuda", dtype=dtype) + reference = torch.cat([baseline(row[:, None]) for row in inputs], dim=1) + run = Path(os.environ["EXP_RUN_DIR"]) + report = dict( + status="running", + scope="trained rank-zero head component; no distributed serving or vLLM/SGLang performance claim", + revision=REVISION, + configuration=vars(args) | {"model_directory": str(args.model_directory)}, + local_weight_shape=list(baseline.weight.shape), + hardware=subprocess.check_output( + ["nvidia-smi", "--query-gpu=index,name,uuid,memory.used,memory.free", "--format=csv,noheader"], text=True + ), + software={name: importlib.metadata.version(name) for name in ("torch", "triton")}, + cuda=torch.version.cuda, + bf16_reduced_precision=torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction, + source_file=verified["files"][filename], + results=[], + ) + path = run / "results.json" + try: + for rows in args.rows: + tokens = inputs[:rows].T + old, new = baseline(tokens), optimized(tokens) + torch.testing.assert_close(new, reference[:, :rows], atol=0, rtol=0) + timings = {} + for name, weight in (("baseline", baseline), ("optimized", optimized)): + timings[name] = [ + triton.testing.do_bench_cudagraph(lambda: weight(tokens), rep=args.rep_ms) * 1000 + for _ in range(args.samples) + ] + item = dict( + rows=rows, + scalar_matches=True, + baseline_max_abs=float((old - reference[:, :rows]).float().abs().max()), + microseconds=timings, + speedup=sorted(timings["baseline"])[args.samples // 2] + / sorted(timings["optimized"])[args.samples // 2], + ) + report["results"].append(item) + path.write_text(json.dumps(report, indent=2)) + print(json.dumps(item), flush=True) + report["status"] = "passed" + except Exception as error: + report.update(status="failed", error_type=type(error).__name__, error=str(error)) + raise + finally: + path.write_text(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/benchmark_ngram_hash.py b/test/benchmark/qwen38/benchmark_ngram_hash.py new file mode 100644 index 0000000000..35042cadcb --- /dev/null +++ b/test/benchmark/qwen38/benchmark_ngram_hash.py @@ -0,0 +1,129 @@ +import argparse +import json +import os + +import torch +import triton +from triton.testing import do_bench_cudagraph + +from lightllm.common.ngram_embedding import NGramHasher + + +def eager(hasher, tokens, history): + batch, length = tokens.shape + full = torch.cat((history, tokens.long()), dim=-1) + positions = torch.arange(full.shape[1], device=full.device, dtype=torch.int64) + eos_positions = torch.where(full == hasher.eos_token_id, positions, -1) + inclusive = torch.cummax(eos_positions, dim=1).values + previous_eos = torch.cat((eos_positions.new_full((batch, 1), -1), inclusive[:, :-1]), dim=1) + relative = positions.unsqueeze(0) - previous_eos - 1 + shifted = [full] + for shift in range(1, hasher.context_len + 1): + source = positions - shift + gathered = full.gather(1, source.clamp_min(0).unsqueeze(0).expand(batch, -1)) + valid = (relative >= shift) & (source.unsqueeze(0) >= 0) + shifted.append(torch.where(valid, gathered, hasher.eos_token_id)) + ids = [] + per_order = hasher.heads // hasher.context_len + for order in range(2, hasher.context_len + 2): + mixed = shifted[0] * hasher.multipliers[0] + for index in range(1, order): + mixed = torch.bitwise_xor(mixed, shifted[index] * hasher.multipliers[index]) + begin = (order - 2) * per_order + sizes = hasher.head_vocab_sizes[begin : begin + per_order] + offsets = hasher.head_offsets[begin : begin + per_order] + ids.append(torch.remainder(mixed.unsqueeze(-1), sizes.view(1, 1, -1)) + offsets.view(1, 1, -1)) + return torch.cat(ids, dim=-1)[:, -length:].reshape(-1, hasher.heads), full[:, -hasher.context_len :].contiguous() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--shapes", nargs="+", default=["1x1", "16x1", "128x1", "1x4096", "16x256"]) + parser.add_argument("--repetitions", type=int, default=3) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Wrap this benchmark with exp -m "purpose"') + torch.manual_seed(37) + + sizes = torch.tensor( + [ + 20000003, + 20000023, + 20000033, + 20000047, + 20000059, + 20000063, + 20000069, + 20000077, + 20000081, + 20000093, + 20000107, + 20000147, + 20000153, + 20000159, + 20000161, + 20000171, + ], + device="cuda", + dtype=torch.int64, + ) + offsets = torch.cat([sizes.new_zeros(1), sizes.cumsum(0)[:-1]]) + multipliers = torch.tensor([23703573157769, 20109073645365, 8052911324071], device="cuda", dtype=torch.int64) + hasher = NGramHasher(multipliers, sizes, offsets, 248044) + print( + json.dumps( + { + "torch": torch.__version__, + "triton": triton.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "ngram_size": 3, + "heads_per_ngram": 8, + "multipliers": multipliers.tolist(), + "head_vocab_sizes": sizes.tolist(), + "eos_token_id": hasher.eos_token_id, + "timing": "CUDA graph, median of interleaved repetitions", + "includes_embedding_lookup": False, + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + } + ), + flush=True, + ) + for shape in args.shapes: + batch, length = map(int, shape.split("x")) + tokens = torch.randint(0, 248320, (batch, length), dtype=torch.int64, device="cuda") + tokens[:, ::31] = hasher.eos_token_id + starts = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * length + history = hasher.empty_history(batch) + history[:, -1] = 77 + + def fused(): + return hasher.hash(tokens.view(-1), starts, history) + + def reference(): + return eager(hasher, tokens, history) + + for actual, expected in zip(fused(), reference()): + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + samples = {"eager": [], "fused": []} + for repeat in range(args.repetitions): + functions = [("eager", reference), ("fused", fused)] + for name, func in functions[:: 1 if repeat % 2 == 0 else -1]: + samples[name].append(do_bench_cudagraph(func, rep=100) * 1000) + medians = {name + "_us": sorted(values)[len(values) // 2] for name, values in samples.items()} + print( + json.dumps( + { + "batch": batch, + "tokens_per_request": length, + **medians, + "speedup": medians["eager_us"] / medians["fused_us"], + "samples_us": samples, + } + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/benchmark_ple.py b/test/benchmark/qwen38/benchmark_ple.py new file mode 100644 index 0000000000..1c98a8481f --- /dev/null +++ b/test/benchmark/qwen38/benchmark_ple.py @@ -0,0 +1,120 @@ +import argparse +import json +import math +import os + +import torch +import torch.nn.functional as F +import triton +from triton.testing import do_bench_cudagraph + +from lightllm.common.ngram_embedding.ple import PLELayer + + +def eager(layer, residual, embeddings, history): + batch, length, width = residual.shape + + def norm(x, weight): + grouped = x.float().unflatten(-1, (layer.branches, -1)) + normalized = grouped * torch.rsqrt(grouped.square().mean(-1, keepdim=True) + layer.eps) + return (normalized.flatten(-2) * (1 + weight.float())).to(x.dtype) + + key = norm(F.linear(embeddings, layer.key_weight), layer.norm_key).unflatten(-1, (layer.branches, -1)) + query = norm(residual, layer.norm_query).unflatten(-1, (layer.branches, -1)) + value = F.linear(embeddings, layer.value_weight) + gate = (key * query).sum(-1, keepdim=True) / math.sqrt(layer.hidden_size) + gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() + gated = (torch.sigmoid(gate) * value.unsqueeze(-2)).flatten(-2) + conv_inputs = norm(gated, layer.norm_conv) + full = torch.cat((history, conv_inputs), dim=1) + conv = F.conv1d(full.transpose(1, 2), layer.conv_weight.unsqueeze(1), dilation=layer.dilation, groups=width) + output = gated + F.silu(conv).transpose(1, 2) + return output, full[:, -layer.context_len :].contiguous() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--shapes", nargs="+", default=["1x1", "16x1", "1x128", "4x128"]) + parser.add_argument("--repetitions", type=int, default=3) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Wrap this benchmark with exp -m "purpose"') + torch.manual_seed(137) + dim, branches, embed_dim = 2560, 4, 2560 + width = dim * branches + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) + + layer = PLELayer( + rand(width, embed_dim) / embed_dim ** 0.5, + rand(dim, embed_dim) / embed_dim ** 0.5, + rand(width) * 0.1, + rand(width) * 0.1, + rand(width) * 0.1, + rand(width, 4) * 0.1, + branches, + 3, + ) + print( + json.dumps( + { + "torch": torch.__version__, + "triton": triton.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "dtype": "bfloat16", + "hidden_size": dim, + "branches": branches, + "embedding_width": embed_dim, + "kernel_size": 4, + "dilation": 3, + "timing": "CUDA graph, median of interleaved repetitions", + "includes_table_lookup": False, + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + } + ), + flush=True, + ) + for shape in args.shapes: + batch, length = map(int, shape.split("x")) + residual = rand(batch, length, width) + embeddings = rand(batch, length, embed_dim) + history = layer.empty_history(batch).normal_() + starts = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * length + + def fused(): + output, conv_inputs = layer.forward( + residual.view(-1, width), embeddings.view(-1, embed_dim), starts, history + ) + return output.view(batch, length, width), layer.history_after(conv_inputs, starts, history) + + def reference(): + return eager(layer, residual, embeddings, history) + + for actual, expected in zip(fused(), reference()): + torch.testing.assert_close( + actual, expected, atol=2 * torch.finfo(residual.dtype).eps, rtol=2 * torch.finfo(residual.dtype).eps + ) + samples = {"eager": [], "fused": []} + for repeat in range(args.repetitions): + functions = [("eager", reference), ("fused", fused)] + for name, func in functions[:: 1 if repeat % 2 == 0 else -1]: + samples[name].append(do_bench_cudagraph(func, rep=100) * 1000) + medians = {name + "_us": sorted(values)[len(values) // 2] for name, values in samples.items()} + print( + json.dumps( + { + "batch": batch, + "tokens_per_request": length, + **medians, + "speedup": medians["eager_us"] / medians["fused_us"], + "samples_us": samples, + } + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/benchmark_qsa.py b/test/benchmark/qwen38/benchmark_qsa.py new file mode 100644 index 0000000000..53eab467dd --- /dev/null +++ b/test/benchmark/qwen38/benchmark_qsa.py @@ -0,0 +1,165 @@ +import argparse +import json +import os + +import torch +import torch.nn.functional as F +import triton +from triton.testing import do_bench_cudagraph + +from lightllm.common.qsa.indexer import QSAIndexer +from lightllm.common.qsa.attention import indexed_attention + + +def eager_select(indexer, query, cache, requests, lengths, table): + blocks = cache.shape[1] + count = lengths // indexer.compress_ratio + + scores = torch.matmul(query.float(), cache[0].float().transpose(-1, -2)) + scores = scores.relu().sum(1) / indexer.head_dim ** 0.5 + scores.masked_fill_(torch.arange(blocks, device=query.device)[None, :] >= count[:, None], -float("inf")) + topk = min(blocks, indexer.token_budget // indexer.compress_ratio) + selected = scores.topk(topk, sorted=False).indices + width = indexer.token_budget + indexer.compress_ratio - 1 + col = torch.arange(width, device=query.device)[None, :].expand(query.shape[0], -1) + chosen = count.clamp_max(indexer.token_budget // indexer.compress_ratio) + full = col < chosen[:, None] * indexer.compress_ratio + rank = (col // indexer.compress_ratio).clamp_max(topk - 1) + position = selected.gather(1, rank) * indexer.compress_ratio + col % indexer.compress_ratio + tail_col = col - chosen[:, None] * indexer.compress_ratio + tail = (tail_col >= 0) & (tail_col < (lengths % indexer.compress_ratio)[:, None]) + position = torch.where(full, position, count[:, None] * indexer.compress_ratio + tail_col) + slots = table[requests.long()[:, None], position.clamp(0, table.shape[1] - 1)] + return slots.masked_fill(~(full | tail), -1) + + +def eager_attention(q, k, v, selected, valid_tokens=2048): + from torch.nn.attention import SDPBackend, sdpa_kernel + + valid_slots = selected[:, :valid_tokens] + gathered_k = k[valid_slots].transpose(1, 2) + gathered_v = v[valid_slots].transpose(1, 2) + with sdpa_kernel(SDPBackend.FLASH_ATTENTION): + return F.scaled_dot_product_attention(q.unsqueeze(2), gathered_k, gathered_v, enable_gqa=True).squeeze(2) + + +def measure(reference, fused, repetitions): + samples = {"torch": [], "kernel": []} + for repeat in range(repetitions): + functions = [("torch", reference), ("kernel", fused)] + for name, function in functions[:: 1 if repeat % 2 == 0 else -1]: + samples[name].append(do_bench_cudagraph(function, rep=100) * 1000) + result = {name + "_us": sorted(values)[len(values) // 2] for name, values in samples.items()} + return {**result, "speedup": result["torch_us"] / result["kernel_us"], "samples_us": samples} + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--contexts", nargs="+", type=int, default=[4099, 65539]) + parser.add_argument("--index-rows", nargs="+", type=int, default=[1, 16]) + parser.add_argument("--attention-rows", nargs="+", type=int, default=[1, 16, 128]) + parser.add_argument("--valid-selections", nargs="+", type=int, default=[2048]) + parser.add_argument("--selection-bounds", action="store_true", help="Use native request visibility bounds") + parser.add_argument("--repetitions", type=int, default=3) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Wrap this benchmark with exp -m "purpose"') + if any(count < 1 or count > 2048 for count in args.valid_selections): + parser.error("Valid selection counts must lie in [1, 2048]") + torch.manual_seed(537) + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) + + indexer = QSAIndexer(rand(640, 2560), rand(128), rand(128)) + print( + json.dumps( + { + "torch": torch.__version__, + "triton": triton.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "dtype": "bfloat16", + "index_heads": 4, + "index_head_dim": 128, + "token_budget": 2048, + "compress_ratio": 4, + "query_heads": 24, + "kv_heads": 2, + "attention_head_dim": 256, + "timing": "CUDA graph, median of interleaved repetitions", + "attention_baseline": "KV gather plus forced Torch FlashAttention, no padding mask", + "selection_baseline": ( + "Torch matmul with shared cached keys, ReLU head reduction, top-k and slot gather" + ), + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "selection_bounds": args.selection_bounds, + } + ), + flush=True, + ) + for context in args.contexts: + cache = rand(1, context // 4, 128) + table = torch.arange(context, device="cuda", dtype=torch.int32).view(1, -1) + for rows in args.index_rows: + query = rand(rows, 4, 128) + requests = torch.zeros(rows, device="cuda", dtype=torch.int32) + lengths = ( + torch.full((rows,), context, device="cuda", dtype=torch.int32) + - torch.arange(rows, device="cuda", dtype=torch.int32) % 7 + ) + + def reference(): + return eager_select(indexer, query, cache, requests, lengths, table) + + def fused(): + return indexer.select(query, cache, requests, lengths, table) + + torch.testing.assert_close(fused().sort(-1).values, reference().sort(-1).values, atol=0, rtol=0) + print( + json.dumps( + { + "component": "cached_block_selection", + "queries": rows, + "context": context, + **measure(reference, fused, args.repetitions), + } + ), + flush=True, + ) + slots = 65539 + key, value = rand(slots, 2, 256), rand(slots, 2, 256) + for rows in args.attention_rows: + query = rand(rows, 24, 256) + choices = torch.stack([torch.randperm(slots, device="cuda")[:2051] for _ in range(rows)]).int() + for count in args.valid_selections: + selected = choices.clone() + selected[:, count:] = -1 + + bound = slots if count == 2048 else count + lengths = torch.full((rows,), bound, device="cuda", dtype=torch.int32) if args.selection_bounds else None + + def reference(): + return eager_attention(query, key, value, selected, count) + + def fused(): + return indexed_attention(query, key, value, selected, selection_lengths=lengths) + + torch.testing.assert_close(fused(), reference(), atol=0.016, rtol=0.016) + print( + json.dumps( + { + "component": "indexed_gqa", + "queries": rows, + "physical_slots": slots, + "selected": count, + "selection_capacity": selected.shape[1], + **measure(reference, fused, args.repetitions), + } + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_aiperf_timing.py b/test/benchmark/qwen38/check_aiperf_timing.py new file mode 100644 index 0000000000..f282df9374 --- /dev/null +++ b/test/benchmark/qwen38/check_aiperf_timing.py @@ -0,0 +1,60 @@ +import argparse +import hashlib +import json +import os +from pathlib import Path + +from test.benchmark.qwen38.measure_aiperf import audit_records, profiling_records +from test.benchmark.service.stream_timing import summarize_streaming_times + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--replay", type=Path, required=True, help="Completed aiperf-replay.json") + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + replay = json.loads(args.replay.read_text()) + if replay["status"] != "aiperf_replay_recorded": + raise ValueError("A completed replay is required") + directory = args.replay.parent + trace_path = directory / "trace.jsonl" + trace = [json.loads(line) for line in trace_path.read_text().splitlines()] + audit = audit_records(directory / "aiperf", trace) + if audit != replay["audit"]: + raise ValueError("Retained replay audit differs from its raw records") + timing = summarize_streaming_times(profiling_records(directory / "aiperf")) + if timing["requests"] != audit["requests"] or timing["empty_text_responses"] != audit["content_errors"]: + raise ValueError("Timing cohort differs from the complete replay") + service = replay["service"] + for key in ("framework", "version", "software", "hardware", "command", "configuration", "model_revision"): + if not service.get(key): + raise ValueError(f"Missing actual service provenance: {key}") + paths = ( + args.replay, + trace_path, + directory / "aiperf/profile_export_raw.jsonl", + Path(__file__), + Path(__file__).parent.parent / "service/stream_timing.py", + ) + report = dict( + status="complete_aiperf_stream_timing_verified", + replay=str(args.replay.resolve()), + service=service, + timing=timing, + aiperf_metrics=audit["metrics"], + hashes={str(path.resolve()): hashlib.sha256(path.read_bytes()).hexdigest() for path in paths}, + comparison_eligible=False, + limitations=( + "Timing only. Complete work and explicit empty-text errors are retained. " + "HTTP completion includes SSE DONE; visible-text boundaries exclude later invisible-token generation. " + "Task quality and matched framework comparisons require separate audits." + ), + ) + run = Path(os.environ["EXP_RUN_DIR"]) + (run / "stream-timing.json").write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(dict(status=report["status"], timing=timing))) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_baseline.py b/test/benchmark/qwen38/check_baseline.py new file mode 100644 index 0000000000..e732a83b9b --- /dev/null +++ b/test/benchmark/qwen38/check_baseline.py @@ -0,0 +1,102 @@ +import argparse +from datetime import datetime, timezone +import hashlib +import importlib.metadata +import inspect +import json +import os +from pathlib import Path +import platform +import subprocess +import sys + + +def check_runtime(framework, model_directory): + report = { + "framework": framework, + "created_at": datetime.now(timezone.utc).isoformat(), + "python": platform.python_version(), + "glibc": platform.libc_ver(), + "executable": sys.executable, + "scope": "Dependency/import/GPU smoke checks only; no serving or performance result", + "software": {}, + "status": "checking", + } + for package in (framework, "torch", "triton", "transformers", "flashinfer-python"): + try: + report["software"][package] = importlib.metadata.version(package) + except importlib.metadata.PackageNotFoundError: + report["software"][package] = None + try: + dependencies = subprocess.run([sys.executable, "-m", "pip", "check"], capture_output=True, text=True) + report["dependencies"] = { + "exit_code": dependencies.returncode, + "output": dependencies.stdout + dependencies.stderr, + } + if dependencies.returncode: + raise RuntimeError("Baseline dependency validation failed") + import torch + + if not torch.cuda.is_available(): + raise RuntimeError("Baseline cannot access CUDA") + report.update(cuda=torch.version.cuda, nccl=torch.cuda.nccl.version(), gpu=torch.cuda.get_device_name()) + a = torch.ones(16, 64, device="cuda", dtype=torch.bfloat16) + b = torch.ones(64, 32, device="cuda", dtype=torch.bfloat16) + torch.testing.assert_close(a @ b, torch.full((16, 32), 64, device="cuda", dtype=torch.bfloat16), atol=0, rtol=0) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = a @ b + a.fill_(2) + graph.replay() + torch.testing.assert_close(captured, torch.full_like(captured, 128), atol=0, rtol=0) + graph.reset() + report["bf16_matmul_and_graph"] = "passed" + architecture = "Qwen4ExpForConditionalGeneration" + if framework == "vllm": + from vllm.config import ModelConfig + from vllm.model_executor.models import ModelRegistry + + config = ModelConfig( + model=str(model_directory.resolve()), skip_tokenizer_init=True, dtype="bfloat16", max_model_len=32 + ) + model_class, resolved_architecture = ModelRegistry.resolve_model_cls(architecture, config) + if resolved_architecture != architecture: + raise RuntimeError("Baseline resolved a different model architecture") + else: + from sglang.srt.models.qwen4_exp import Qwen4ExpForConditionalGeneration + + model_class = Qwen4ExpForConditionalGeneration + source = Path(inspect.getfile(model_class)) + report["model"] = { + "architecture": architecture, + "module": model_class.__module__, + "source": str(source), + "sha256": hashlib.sha256(source.read_bytes()).hexdigest(), + } + report["status"] = "passed" + except Exception as error: + report.update(status="failed", error_type=type(error).__name__, error=str(error)) + return report + + +def main(): + if not os.environ.get("EXP_ACTIVE"): + raise RuntimeError("Run baseline checks through exp") + parser = argparse.ArgumentParser() + parser.add_argument("--framework", choices=("vllm", "sglang"), required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--model-directory", type=Path, default=Path("artifacts/references/qwen4_exp/layer0")) + args = parser.parse_args() + result = check_runtime(args.framework, args.model_directory) + payload = json.dumps(result, indent=2) + "\n" + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(payload) + (Path(os.environ["EXP_RUN_DIR"]) / "baseline-readiness.json").write_text(payload) + print(payload, flush=True) + if result["status"] != "passed": + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_cold_streams.py b/test/benchmark/qwen38/check_cold_streams.py new file mode 100644 index 0000000000..2c5da4f2d9 --- /dev/null +++ b/test/benchmark/qwen38/check_cold_streams.py @@ -0,0 +1,159 @@ +import argparse +import hashlib +import json +import os +from pathlib import Path +import secrets +import shutil + +from test.benchmark.qwen38.check_generation import REVISION, cases_for, verified_checkpoint +from test.benchmark.service.stream_cohort import StreamingWorkload, run_cohort, summarize_cohort +from test.benchmark.service.token_stream import request_stream + + +def workloads_from_reference(reference, expected_cases): + if ( + reference.get("status") not in ("passed", "recorded") + or reference.get("revision") != REVISION + or reference.get("cases") != expected_cases + ): + raise ValueError("Require a qualified generation reference with the exact pinned inputs") + length = reference.get("output_length") + if type(length) is not int or length < 1: + raise ValueError("Reference output length must be a positive integer") + results = reference.get("results", {}) + if set(results) != set(expected_cases): + raise ValueError("Reference must contain every declared generation case") + workloads = [] + for name, prompt in expected_cases.items(): + ids = results[name].get("token_ids") + if not isinstance(ids, list) or len(ids) != length or any(type(token) is not int or token < 0 for token in ids): + raise ValueError("Reference must contain the complete actual output IDs") + if name.startswith("repeat_"): + original = name.removeprefix("repeat_") + if ( + original not in expected_cases + or prompt != expected_cases[original] + or ids != results[original]["token_ids"] + ): + raise ValueError("Repeated reference prompts and outputs must match") + else: + workloads.append(StreamingWorkload(name, prompt, ids)) + if not workloads: + raise ValueError("Require at least one non-repeated workload") + return workloads + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--repeats", type=int, default=3) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + if not 1 <= args.repeats <= 20: + parser.error("Repeats must be between 1 and 20") + run = Path(os.environ["EXP_RUN_DIR"]) + shutil.copyfile(__file__, run / "cold-stream-cohorts.py") + verified_checkpoint(args.model_directory) + service = json.loads(args.service_metadata.read_text()) + if service.get("framework") not in ("lightllm", "vllm", "sglang") or any( + not service.get(key) for key in ("version", "command", "hardware", "configuration") + ): + raise ValueError("Record actual service framework, version, command, hardware and configuration") + reference = json.loads(args.reference.read_text()) + lengths = [int(name.removeprefix("long_")) for name in reference.get("cases", {}) if name.startswith("long_")] + workloads = workloads_from_reference(reference, cases_for(args.model_directory, lengths)) + manifest = [dict(name=work.name, prompt_ids=work.prompt_ids, expected_ids=work.expected_ids) for work in workloads] + report = dict( + status="running", + comparison_eligible=False, + scope=__doc__, + reason="Shared hardware diagnostic; exact-output and explicit zero-cache requirements remain mandatory", + service=service, + reference=str(args.reference), + revision=REVISION, + manifest=manifest, + manifest_sha256=hashlib.sha256(json.dumps(manifest, sort_keys=True).encode()).hexdigest(), + cache_policy="native request cache bypass or independent random baseline cache salt on every request", + repeats=args.repeats, + cohorts=[], + ) + output = run / "cold-cohort.json" + + def save(): + output.write_text(json.dumps(report, indent=2)) + + def request(workload, response_path): + options = ( + dict(disable_prompt_cache=True) + if service["framework"] == "lightllm" + else dict(cache_salt=secrets.token_urlsafe(32)) + ) + return request_stream( + args.base_url, + service["framework"], + str(args.model_directory.resolve()), + list(workload.prompt_ids), + len(workload.expected_ids), + 600, + response_path, + **options, + ) + + try: + save() + for repeat in range(args.repeats): + order = [1, 2, 4, 8] + order = order[repeat % 4 :] + order[: repeat % 4] + sequence = workloads[repeat % len(workloads) :] + workloads[: repeat % len(workloads)] + for concurrency in order: + result = run_cohort( + sequence * 4, concurrency, request, run / f"repeat-{repeat}/concurrency-{concurrency}" + ) + result["repeat"] = repeat + transport_complete = all( + row.get("result") is not None + and ( + row["status"] == "passed" + or row.get("error") == "Output IDs differ from the qualified reference" + ) + for row in result["records"] + ) + cached = [ + row.get("result", {}).get("metrics", {}).get("cached_input_tokens") for row in result["records"] + ] + cold = all(type(count) is int and count == 0 for count in cached) + result["zero_cache_observed"] = cold + if not cold: + result.update( + status="failed", + metrics=None, + cache_gate_error="Every request must report exactly zero cached input tokens", + ) + result["unqualified_diagnostic_metrics"] = ( + summarize_cohort([row["result"] for row in result["records"]]) if transport_complete else None + ) + report["cohorts"].append(result) + save() + print( + json.dumps( + dict(repeat=repeat, concurrency=concurrency, status=result["status"], zero_cache_observed=cold) + ), + flush=True, + ) + if result["status"] != "passed": + raise ValueError("Cold request quality, timing or cache gate failed; qualified metrics unavailable") + report["status"] = "passed" + except Exception as error: + report.update(status="failed", error_type=type(error).__name__, error=str(error)) + raise + finally: + save() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_cpu_cache.py b/test/benchmark/qwen38/check_cpu_cache.py new file mode 100644 index 0000000000..e5aa24f4dd --- /dev/null +++ b/test/benchmark/qwen38/check_cpu_cache.py @@ -0,0 +1,146 @@ +import argparse +from concurrent.futures import ThreadPoolExecutor +import json +import os +from pathlib import Path + +from check_generation import REVISION, verified_checkpoint +from check_prefix_eviction import cache_hit, eviction_cases, validate_eviction +from check_scheduler import generate + + +def validate_cpu_restore(events, ranks, target_token, initial, restored, minimum_prefix=2048): + if type(ranks) is not int or ranks < 1 or not initial < restored[0] < restored[1]: + raise ValueError("Require valid rank count and ordered restore interval") + fields = {"qsa_raw_keys", "qsa_pooled_keys", "qsa_key_cos", "qsa_key_sin"} + proof = [] + for rank in range(ranks): + local = [event for event in events if event.get("rank") == rank and event.get("first_token") == target_token] + offloads = [ + event + for event in local + if event["event"] == "cpu_offload_queued" + and initial <= event["monotonic"] < restored[0] + and type(event.get("pages")) is int + and event["pages"] > 0 + ] + loads = [ + event + for event in local + if event["event"] == "cpu_load_queued" + and restored[0] <= event["monotonic"] <= restored[1] + and type(event.get("tokens")) is int + and event["tokens"] >= 128 + and type(event.get("kv_length")) is int + and event["kv_length"] >= minimum_prefix + and type(event.get("request_id")) is int + and event["request_id"] >= 0 + and event.get("pinned") is True + ] + complete = [] + for load in loads: + extra = [ + event + for event in local + if event["event"] == "cpu_extra_load_queued" + and event.get("request_id") == load.get("request_id") + and restored[0] <= event["monotonic"] <= load["monotonic"] + and fields <= set(event.get("slot_fields", [])) + and event.get("lexical_layers") + ] + if extra: + complete.append(load) + if not offloads or not complete: + raise ValueError(f"No complete target CPU offload/restore evidence on TP rank {rank}") + proof.append(dict(rank=rank, offloads=len(offloads), restores=len(complete))) + return proof + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--cache-events", type=Path, required=True) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--timeout", type=float, default=600) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error("Run through exp") + verified_checkpoint(args.model_directory) + service = json.loads(args.service_metadata.read_text()) + if ( + service.get("framework") != "lightllm" + or any(not service.get(key) for key in ("version", "command", "hardware", "configuration")) + or service["configuration"].get("cpu_cache") is not True + ): + raise ValueError("Require complete actual native CPU-cache service provenance") + cases = eviction_cases(args.model_directory) + reference = json.loads(args.reference.read_text()) + if ( + reference.get("status") not in ("recorded", "passed") + or reference["revision"] != REVISION + or reference["cases"] != cases + ): + raise ValueError("Require a passing reference with identical checkpoint and pressure inputs") + run = Path(os.environ["EXP_RUN_DIR"]) + report = dict(status="running", scope=__doc__, service=service, revision=REVISION, cases=cases, requests={}) + path = run / "cpu-cache.json" + + def request(label, case): + result = generate(args.base_url, case, [], args.timeout, run / f"{label}.json") + if result["token_ids"] != reference["requests"][label]["token_ids"]: + raise ValueError(f"CPU-cache output differs from qualified reference for {label}") + cache_hit(result) + report["requests"][label] = result + path.write_text(json.dumps(report, indent=2)) + return result + + try: + target = cases["target"] + initial = request("initial", target) + warm = request("warm", target) + if cache_hit(initial) != 0 or cache_hit(warm) < 2048: + raise ValueError("Require fresh target followed by actual warm-prefix reuse") + for name, case in cases.items(): + if name.startswith("small_pressure_"): + request(name, case) + request("after_small_pressure", target) + for name, case in cases.items(): + if name.startswith("pressure_"): + request(name, case) + after = request("after_pressure", target) + if cache_hit(after) < 2048: + raise ValueError("Evicted target did not restore a long cached prefix") + request("rewarm", target) + + with ThreadPoolExecutor(max_workers=4) as executor: + results = list( + executor.map( + lambda index: generate(args.base_url, target, [], args.timeout, run / f"concurrent_{index}.json"), + range(4), + ) + ) + for index, result in enumerate(results): + label = f"concurrent_{index}" + if result["token_ids"] != reference["requests"][label]["token_ids"]: + raise ValueError("Concurrent restored outputs differ from qualified reference") + report["requests"][label] = result + events = [json.loads(line) for line in args.cache_events.read_text().splitlines() if line] + ranks = service["configuration"]["tp"] + report["eviction_proof"] = validate_eviction( + events, ranks, target["inputs"][0], warm["ended"], after["started"] + ) + report["cpu_restore_proof"] = validate_cpu_restore( + events, ranks, target["inputs"][0], initial["started"], (after["started"], after["ended"]) + ) + report["status"] = "passed" + except Exception as error: + report.update(status="failed", error_type=type(error).__name__, error=str(error)) + raise + finally: + path.write_text(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_generation.py b/test/benchmark/qwen38/check_generation.py new file mode 100644 index 0000000000..b0632992b0 --- /dev/null +++ b/test/benchmark/qwen38/check_generation.py @@ -0,0 +1,228 @@ +import argparse +import hashlib +import json +import os +from pathlib import Path +import secrets + +import requests +from tokenizers import Tokenizer + +REVISION = "de4b8e4d43b917e7706784d8bb445c9af86a3540" +TOKENIZER_SHA256 = "0997f410c57a1f4e53b09e4be8f4a172d90edd9564368fb0847030937229b9f3" + + +def verified_checkpoint(root): + acquisition = json.loads((root / "acquisition.json").read_text()) + verified = json.loads((root / "verified_files.json").read_text()) + if acquisition["status"] != "verified" or acquisition["revision"] != REVISION or verified["revision"] != REVISION: + raise ValueError("The complete pinned checkpoint must finish verification before this check") + index = json.loads((root / "model.safetensors.index.json").read_text()) + required = set(index["weight_map"].values()) | {"config.json", "model.safetensors.index.json"} + for name in required: + metadata = verified["files"][name] + if (root / name).stat().st_size != metadata["size"]: + raise ValueError(f"Checkpoint file size changed after verification: {name}") + return verified + + +def cases_for(root, long_lengths=(2049, 4063)): + if not long_lengths or any(type(length) is not int or length < 1 for length in long_lengths): + raise ValueError("Long prompt lengths must be positive integers") + if len(set(long_lengths)) != len(long_lengths): + raise ValueError("Long prompt lengths must be distinct") + path = root / "tokenizer.json" + if hashlib.sha256(path.read_bytes()).hexdigest() != TOKENIZER_SHA256: + raise ValueError("Tokenizer does not match the pinned checkpoint") + tokenizer = Tokenizer.from_file(str(path)) + prompts = { + "english": "The capital of France is", + "chinese": "中国的首都是", + "code_unicode": "def add(a, b):\n # 求和 🐍\n return", + } + cases = {name: tokenizer.encode(prompt, add_special_tokens=False).ids for name, prompt in prompts.items()} + seed = tokenizer.encode("A short factual note about language models and their context. ").ids + for length in long_lengths: + cases[f"long_{length}"] = (seed * (length // len(seed) + 1))[:length] + cases[f"repeat_long_{length}"] = list(cases[f"long_{length}"]) + return cases + + +def generate( + client, base_url, framework, model, tokens, output_length, timeout, response_path=None, *, cold_cache=False +): + if framework == "lightllm": + endpoint = "/generate" + payload = { + "inputs": tokens, + "parameters": { + "do_sample": False, + "ignore_eos": True, + "max_new_tokens": output_length, + "return_details": True, + }, + } + if cold_cache: + payload["parameters"]["disable_prompt_cache"] = True + else: + endpoint = "/v1/completions" + payload = { + "model": model, + "prompt": tokens, + "max_tokens": output_length, + "temperature": 0, + "seed": 3107, + "ignore_eos": True, + "return_token_ids": True, + } + if cold_cache: + payload["cache_salt"] = secrets.token_urlsafe(32) + response = client.post(base_url.rstrip("/") + endpoint, json=payload, timeout=timeout) + if response_path is not None: + response_path.write_text(response.text) + response.raise_for_status() + body = response.json() + if framework == "lightllm": + ids = [token["id"] for token in body["tokens"]] + prompt_length = body["prompt_tokens"] + else: + ids = body["choices"][0]["token_ids"] + prompt_length = body["usage"]["prompt_tokens"] + if not isinstance(ids, list) or len(ids) != output_length or not all(type(token) is int for token in ids): + received = len(ids) if isinstance(ids, list) else type(ids).__name__ + raise ValueError(f"Expected {output_length} integer token IDs, received {received}") + if prompt_length != len(tokens): + raise ValueError("Service prompt token count does not match the shared tokenized input") + if cold_cache: + cached = ( + [token.get("prompt_cache_len") for token in body["tokens"]] + if framework == "lightllm" + else [(body["usage"].get("prompt_tokens_details") or {}).get("cached_tokens")] + ) + if not cached or any(type(count) is not int or count != 0 for count in cached): + raise ValueError("Cold generation must explicitly report zero cached input tokens") + return {"token_ids": ids, "response": body} + + +def main(): + if not os.environ.get("EXP_ACTIVE"): + raise RuntimeError("Run through exp") + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--framework", choices=("lightllm", "vllm", "sglang"), required=True) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--reference", type=Path) + parser.add_argument("--output-length", type=int, default=32) + parser.add_argument("--long-prompt-lengths", type=int, nargs="+", default=[2049, 4063]) + parser.add_argument("--timeout", type=float, default=600) + parser.add_argument( + "--cold-cache", action="store_true", help="Isolate each request and require zero reported cache hits" + ) + parser.add_argument( + "--record-observations", + action="store_true", + help="Retain baseline repeat failures as observations, without qualifying a reference", + ) + args = parser.parse_args() + if args.record_observations and (args.framework == "lightllm" or args.reference or not args.cold_cache): + parser.error("Baseline observations require cold requests and no reference; native qualification stays strict") + if args.output_length < 1: + parser.error("Output length must be positive") + if any(length < 1 for length in args.long_prompt_lengths) or len(set(args.long_prompt_lengths)) != len( + args.long_prompt_lengths + ): + parser.error("Long prompt lengths must be distinct positive integers") + run = Path(os.environ["EXP_RUN_DIR"]) + metadata = json.loads(args.service_metadata.read_text()) + for key in ("framework", "version", "command", "hardware", "configuration"): + if not metadata.get(key): + raise ValueError(f"Service metadata must record {key}") + if metadata["framework"] != args.framework: + raise ValueError("Service metadata names a different framework") + (run / "verified-files.json").write_text(json.dumps(verified_checkpoint(args.model_directory), indent=2)) + cases = cases_for(args.model_directory, args.long_prompt_lengths) + reference = json.loads(args.reference.read_text()) if args.reference else None + if reference and ( + reference["status"] not in ("passed", "recorded") + or reference["revision"] != REVISION + or reference["cases"] != cases + or reference["output_length"] != args.output_length + or reference.get("cache_policy", "service_default") != ("cold" if args.cold_cache else "service_default") + ): + raise ValueError("Reference must use the same revision, exact inputs and output length") + report = { + "status": "running", + "revision": REVISION, + "tokenizer_sha256": TOKENIZER_SHA256, + "service": metadata, + "base_url": args.base_url, + "cases": cases, + "output_length": args.output_length, + "long_prompt_lengths": args.long_prompt_lengths, + "reference": str(args.reference) if args.reference else None, + "cache_policy": "cold" if args.cold_cache else "service_default", + "results": {}, + } + output = run / "generation.json" + client = requests.Session() + client.trust_env = False + active_case = None + try: + for name, tokens in cases.items(): + active_case = name + result = generate( + client, + args.base_url, + args.framework, + str(args.model_directory.resolve()), + tokens, + args.output_length, + args.timeout, + response_path=run / f"response-{name}.txt", + cold_cache=args.cold_cache, + ) + report["results"][name] = result + result["reference_matches"] = ( + result["token_ids"] == reference["results"][name]["token_ids"] if reference else None + ) + output.write_text(json.dumps(report, indent=2)) + print( + json.dumps( + {"case": name, "token_ids": result["token_ids"], "reference_matches": result["reference_matches"]} + ), + flush=True, + ) + active_case = None + results = report["results"] + repeat_matches = all( + results[f"long_{length}"]["token_ids"] == results[f"repeat_long_{length}"]["token_ids"] + for length in args.long_prompt_lengths + ) + report["repeat_consistency"] = dict( + status="passed" if repeat_matches else "failed", + mismatches=[ + length + for length in args.long_prompt_lengths + if results[f"long_{length}"]["token_ids"] != results[f"repeat_long_{length}"]["token_ids"] + ], + ) + if args.record_observations: + report["status"] = "observations_recorded" + report[ + "scope" + ] = "Baseline token/count observations; repeat failures remain explicit, no exact-output qualification" + return + if not repeat_matches or any(item["reference_matches"] is False for item in results.values()): + raise ValueError("Greedy token IDs differ across repeated prompts or reference services") + report["status"] = "passed" if reference else "recorded" + except Exception as error: + report.update(status="failed", failed_case=active_case, error_type=type(error).__name__, error=str(error)) + raise + finally: + output.write_text(json.dumps(report, indent=2)) + client.close() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_math_quality.py b/test/benchmark/qwen38/check_math_quality.py new file mode 100644 index 0000000000..be9652f768 --- /dev/null +++ b/test/benchmark/qwen38/check_math_quality.py @@ -0,0 +1,129 @@ +import argparse +import hashlib +import json +import os +from pathlib import Path + +import requests + +from lightllm.server.tokenizer import get_tokenizer +from test.benchmark.qwen38.check_generation import REVISION, TOKENIZER_SHA256, verified_checkpoint +from test.benchmark.qwen38.math_quality import ( + DATASET_REVISION, + DATASET_SHA256, + load_examples, + prepare_cases, + score_answer, + select_examples, +) +from test.benchmark.service.completion_runner import CompletionClient + + +def verified_workload(manifest, root, dataset_path, tokenizer): + if ( + manifest["status"] != "prepared" + or manifest["model_revision"] != REVISION + or manifest["dataset"]["revision"] != DATASET_REVISION + or manifest["dataset"]["sha256"] != DATASET_SHA256 + or manifest["dataset"]["split"] != "test" + or manifest["decoding"] != dict(temperature=0, ignore_eos=False, enable_thinking=False) + or manifest["tokenizer_files"].get("tokenizer.json") != TOKENIZER_SHA256 + ): + raise ValueError("Workload must use pinned model, dataset and greedy non-thinking decoding") + for name in ("tokenizer.json", "tokenizer_config.json", "chat_template.jinja"): + if hashlib.sha256((root / name).read_bytes()).hexdigest() != manifest["tokenizer_files"].get(name): + raise ValueError("Tokenizer files differ from workload preparation") + examples = load_examples(dataset_path) + if len(examples) != manifest["dataset"]["records"]: + raise ValueError("Dataset record count differs") + cases = prepare_cases( + select_examples(examples, manifest["sample_count"], manifest["sample_seed"]), + tokenizer, + manifest["max_output_tokens"], + manifest["context_limit"], + ) + if cases != manifest["cases"]: + raise ValueError("Workload prompts or gold labels differ from pinned preparation") + return cases + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--framework", choices=("lightllm", "vllm", "sglang"), required=True) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--manifest", type=Path, required=True) + parser.add_argument("--dataset", type=Path, required=True) + parser.add_argument("--timeout", type=float, default=600) + parser.add_argument("--concurrency", type=int, default=1) + args = parser.parse_args() + if not 1 <= args.concurrency <= 8: + parser.error("Concurrency must be between one and eight") + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + run = Path(os.environ["EXP_RUN_DIR"]) + verified_checkpoint(args.model_directory) + metadata = json.loads(args.service_metadata.read_text()) + if metadata.get("framework") != args.framework or any( + not metadata.get(key) for key in ("version", "command", "hardware", "configuration") + ): + raise ValueError("Record actual service framework, version, command, hardware and configuration") + manifest = json.loads(args.manifest.read_text()) + tokenizer = get_tokenizer(str(args.model_directory)).tokenizer + cases = verified_workload(manifest, args.model_directory, args.dataset, tokenizer) + config = json.loads((args.model_directory / "generation_config.json").read_text()) + eos = config["eos_token_id"] + eos_ids = sorted(set(eos if isinstance(eos, list) else [eos]) | {tokenizer.eos_token_id}) + report = dict( + status="running", + scope="Pinned non-thinking scored math sample; no performance or broad accuracy claim", + service=metadata, + manifest=manifest, + manifest_sha256=hashlib.sha256(json.dumps(manifest, sort_keys=True).encode()).hexdigest(), + eos_ids=eos_ids, + concurrency=args.concurrency, + results=[], + ) + output = run / "math-quality.json" + client = CompletionClient( + args.framework, + str(args.model_directory.resolve()), + args.base_url, + args.timeout, + eos_ids, + run, + concurrency=args.concurrency, + session_factory=requests.Session, + ) + active_case = None + try: + responses = iter(client.results(cases)) + for case in cases: + active_case = case["dataset_index"] + result = next(responses) + text = tokenizer.decode(result["token_ids"], skip_special_tokens=True) + score = score_answer(text, case["expected_number"]) + score["correct"] = score["correct"] and not result["truncated"] + result.update(dataset_index=active_case, text=text, **score) + report["results"].append(result) + output.write_text(json.dumps(report, indent=2)) + print(json.dumps({key: result[key] for key in ("dataset_index", "correct", "truncated")}), flush=True) + report.update( + status="recorded", + scored=len(cases), + correct=sum(row["correct"] for row in report["results"]), + truncated=sum(row["truncated"] for row in report["results"]), + invalid_final_answer=sum(row["predicted"] is None for row in report["results"]), + ) + report["accuracy"] = report["correct"] / report["scored"] + except Exception as error: + report.update(status="failed", failed_case=active_case, error_type=type(error).__name__, error=str(error)) + raise + finally: + output.write_text(json.dumps(report, indent=2)) + client.close() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_multimodal.py b/test/benchmark/qwen38/check_multimodal.py new file mode 100644 index 0000000000..9a2bacd63d --- /dev/null +++ b/test/benchmark/qwen38/check_multimodal.py @@ -0,0 +1,291 @@ +import argparse +import base64 +from concurrent.futures import ThreadPoolExecutor, as_completed +import hashlib +from io import BytesIO +import json +import os +from pathlib import Path +import re +from types import SimpleNamespace + +from PIL import Image +import requests + +from check_generation import REVISION, verified_checkpoint +from check_scheduler import validate_generation + + +def image_fixtures(directory, pressure_count): + directory.mkdir(parents=True, exist_ok=True) + definitions = {"red": ((255, 0, 0), (256, 256)), "blue": ((0, 0, 255), (256, 256))} + definitions["rect_blue"] = ((0, 0, 255), (384, 128)) + definitions.update({f"pressure_{i}": ((255, i + 1, 0), (256, 256)) for i in range(pressure_count)}) + payloads, metadata = {}, {} + for name, (color, size) in definitions.items(): + buffer = BytesIO() + Image.new("RGB", size, color).save(buffer, format="PNG") + data = buffer.getvalue() + (directory / f"{name}.png").write_bytes(data) + payloads[name] = {"type": "base64", "data": base64.b64encode(data).decode("ascii")} + metadata[name] = dict(sha256=hashlib.sha256(data).hexdigest(), width=size[0], height=size[1]) + if len({item["sha256"] for item in metadata.values()}) != len(metadata): + raise ValueError("Cache pressure requires distinct image bytes") + return payloads, metadata + + +def prepare_cases(tokenizer, images): + definitions = { + "red": (["red"], ["red"]), + "blue": (["blue"], ["blue"]), + "rect_blue": (["rect_blue"], ["blue"]), + "pair": (["red", "blue"], ["red", "blue"]), + "reversed_pair": (["blue", "red"], ["blue", "red"]), + } + definitions.update({name: ([name], []) for name in images if name.startswith("pressure_")}) + cases = {} + for name, (names, expected) in definitions.items(): + question = ( + "Name the colors of these images in order. Reply with only the two English color words, " + "separated by a comma." + if len(names) > 1 + else "What color is the image? Reply with one English word." + ) + prompt = tokenizer.tokenizer.apply_chat_template( + [{"role": "user", "content": [{"type": "image"} for _ in names] + [{"type": "text", "text": question}]}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + items = [] + for index, image in enumerate(names): + item = SimpleNamespace( + image_w=images[image]["width"], image_h=images[image]["height"], token_id=100000000 + index * 1024 + ) + item.token_num = tokenizer.get_image_token_length(item) + items.append(item) + expanded = tokenizer.encode(prompt, SimpleNamespace(images=items, audios=[])) + cases[name] = dict( + inputs=prompt, + images=names, + expected_words=expected, + prompt_tokens=len(expanded), + image_tokens=sum(item.token_num for item in items), + max_new_tokens=1 if not expected else 64, + ignore_eos=not expected, + ) + return cases + + +def validate_image_result(body, case, eos_ids, tokenizer): + ids = validate_generation(body, case["prompt_tokens"], case["max_new_tokens"], case["ignore_eos"], eos_ids) + usage = body["input_usage"] + if ( + usage["input_image_tokens"] != case["image_tokens"] + or usage["input_audio_tokens"] != 0 + or usage["input_text_tokens"] != case["prompt_tokens"] - case["image_tokens"] + ): + raise ValueError("Image/text accounting differs from actual expanded input IDs") + answer = tokenizer.decode([token for token in ids if token not in eos_ids], skip_special_tokens=True).strip() + if case["expected_words"] and re.findall(r"[a-z]+", answer.lower()) != case["expected_words"]: + raise ValueError(f"Image color/order response differs from the fixture: {answer!r}") + return dict(token_ids=ids, answer=answer, response=body) + + +def verify_reallocation(events, image_md5): + allocations = [ + item + for event in events + if event["event"] == "alloc" + for item in event["records"] + if item["md5"].startswith(image_md5 + "_") + ] + if not allocations: + raise ValueError("No actual image-cache allocation was observed") + first, last = allocations[0], allocations[-1] + evicted = {event["record"]["id"] for event in events if event["event"] == "evicted"} + embedded = {item["id"] for event in events if event["event"] == "embedded" for item in event["records"]} + if first["id"] == last["id"] or first["token_id"] == last["token_id"] or first["id"] not in evicted: + raise ValueError("The original image was not demonstrably evicted and assigned new virtual IDs") + if not {first["id"], last["id"]}.issubset(embedded): + raise ValueError("Both original and reallocated images require actual embedding-completion evidence") + return dict(original=first, reallocated=last, original_evicted=True, both_embedded=True) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--reference", type=Path) + parser.add_argument("--cache-events", type=Path) + parser.add_argument("--prepare-only", action="store_true") + parser.add_argument("--timeout", type=float, default=600) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error("Run through exp") + verified_checkpoint(args.model_directory) + service = json.loads(args.service_metadata.read_text()) + if service.get("framework") != "lightllm" or any( + not service.get(key) for key in ("version", "command", "hardware", "configuration") + ): + raise ValueError("Require actual native service provenance") + command = service["command"] + capacity = int(command[command.index("--cache_capacity") + 1]) + if not 3 <= capacity <= 16: + raise ValueError("This bounded two-image workload requires explicit cache capacity between 3 and 16") + run = Path(os.environ["EXP_RUN_DIR"]) + payloads, images = image_fixtures(run / "images", capacity + 1) + from lightllm.server.tokenizer import get_tokenizer + + tokenizer = get_tokenizer(str(args.model_directory)) + cases = prepare_cases(tokenizer, images) + source = { + name: hashlib.sha256((args.model_directory / name).read_bytes()).hexdigest() + for name in ( + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "preprocessor_config.json", + "generation_config.json", + ) + } + generation_eos = json.loads((args.model_directory / "generation_config.json").read_text())["eos_token_id"] + eos_ids = sorted( + set(generation_eos if isinstance(generation_eos, list) else [generation_eos]) | {tokenizer.eos_token_id} + ) + reference = json.loads(args.reference.read_text()) if args.reference else None + if reference and ( + reference["status"] not in ("passed", "recorded") + or reference["revision"] != REVISION + or reference["source"] != source + or reference["cases"] != cases + or reference["images"] != images + or reference["eos_ids"] != eos_ids + ): + raise ValueError("Require a passing reference with identical images, prompt template, processor and EOS") + report = dict( + status="running", + scope=__doc__, + revision=REVISION, + service=service, + source=source, + reference=str(args.reference) if args.reference else None, + images=images, + cases=cases, + eos_ids=eos_ids, + cache_capacity=capacity, + serial={}, + repeats=[], + pressure=[], + waves=[], + ) + path = run / "multimodal.json" + active = None + + def save(): + path.write_text(json.dumps(report, indent=2)) + + if args.prepare_only: + report["status"] = "prepared" + save() + return + + def request(name, label, cold=False): + case = cases[name] + params = dict( + do_sample=False, + return_details=True, + print_eos_token=True, + disable_prompt_cache=cold, + max_new_tokens=case["max_new_tokens"], + ignore_eos=case["ignore_eos"], + ) + with requests.Session() as client: + client.trust_env = False + response = client.post( + args.base_url.rstrip("/") + "/generate", + json=dict( + inputs=case["inputs"], + parameters=params, + multimodal_params=dict(images=[payloads[n] for n in case["images"]]), + ), + timeout=args.timeout, + ) + (run / f"{label}.txt").write_text(response.text) + response.raise_for_status() + return validate_image_result(response.json(), case, eos_ids, tokenizer.tokenizer) + + def compare(name, result): + if result["token_ids"] != report["serial"][name]["token_ids"]: + raise ValueError(f"Repeated/concurrent image IDs differ: {name}") + + try: + save() + for name in ("red", "blue", "rect_blue", "pair", "reversed_pair"): + active = {"phase": "serial", "case": name} + case = cases[name] + with requests.Session() as client: + client.trust_env = False + response = client.post( + args.base_url.rstrip("/") + "/tokens", + json=dict( + text=case["inputs"], multimodal_params=dict(images=[payloads[n] for n in case["images"]]) + ), + timeout=args.timeout, + ) + (run / f"count-{name}.txt").write_text(response.text) + response.raise_for_status() + if response.json()["ntokens"] != case["prompt_tokens"]: + raise ValueError("Public image token count differs from actual expansion") + result = request(name, f"serial-{name}") + report["serial"][name] = result + save() + if reference and result["token_ids"] != reference["serial"][name]["token_ids"]: + raise ValueError(f"Ordinary/MTP image output IDs differ: {name}") + for cold in (False, True): + for name in ("red", "rect_blue", "pair"): + active = dict(phase="repeat", case=name, cold=cold) + result = request(name, f"repeat-{int(cold)}-{name}", cold) + report["repeats"].append(dict(case=name, cold=cold, result=result)) + save() + compare(name, result) + for index in range(capacity + 1): + name = f"pressure_{index}" + active = dict(phase="cache_pressure", case=name) + report["pressure"].append(request(name, name)) + save() + active = dict(phase="after_cache_pressure", case="red") + report["after_cache_pressure"] = request("red", "after-cache-pressure") + save() + compare("red", report["after_cache_pressure"]) + for concurrency in (2, 4): + active = dict(phase="concurrent", concurrency=concurrency) + wave = dict(concurrency=concurrency, results={}) + report["waves"].append(wave) + with ThreadPoolExecutor(max_workers=concurrency) as executor: + pending = { + executor.submit(request, name, f"wave-{concurrency}-{index}"): (index, name) + for index, name in enumerate(["red", "blue", "pair", "reversed_pair"] * 2) + } + for future in as_completed(pending): + index, name = pending[future] + result = future.result() + wave["results"][str(index)] = dict(case=name, result=result) + save() + compare(name, result) + if args.cache_events: + events = [json.loads(line) for line in args.cache_events.read_text().splitlines()] + digest = hashlib.md5((run / "images/red.png").read_bytes()).hexdigest() + report["cache_reallocation"] = verify_reallocation(events, digest) + report["cache_events"] = str(args.cache_events) + report["status"] = "passed" if reference else "recorded" + except Exception as error: + report.update(status="failed", failed_request=active, error_type=type(error).__name__, error=str(error)) + raise + finally: + save() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_prefix_eviction.py b/test/benchmark/qwen38/check_prefix_eviction.py new file mode 100644 index 0000000000..e604bfbaf4 --- /dev/null +++ b/test/benchmark/qwen38/check_prefix_eviction.py @@ -0,0 +1,174 @@ +import argparse +from concurrent.futures import ThreadPoolExecutor +import json +import os +from pathlib import Path + +from check_generation import REVISION, cases_for, verified_checkpoint +from check_scheduler import generate + + +def eviction_cases(root, pressure_count=12, prompt_length=2369): + if pressure_count < 2 or not 2049 <= prompt_length <= 4063: + raise ValueError("Eviction workload requires distinct long prefixes") + seed = cases_for(root)["long_4063"] + cases = { + "target" + if index == 0 + else f"pressure_{index - 1}": dict( + inputs=[1000 + index, *seed[: prompt_length - 1]], + max_new_tokens=32 if index == 0 else 8, + ignore_eos=True, + ) + for index in range(pressure_count + 1) + } + cases.update( + { + f"small_pressure_{index}": dict(inputs=[2000 + index, *seed[:64]], max_new_tokens=8, ignore_eos=True) + for index in range(20) + } + ) + return cases + + +def cache_hit(result): + value = result["response"]["tokens"][0]["prompt_cache_len"] + if type(value) is not int or not 0 <= value < result["response"]["prompt_tokens"]: + raise ValueError("Invalid actual prompt-cache hit length") + return value + + +def validate_eviction(events, ranks, target_token, started, ended): + if type(ranks) is not int or ranks < 1 or not started < ended: + raise ValueError("Require valid TP ranks and pressure interval") + during = [event for event in events if started <= event["monotonic"] <= ended] + proof = [] + for rank in range(ranks): + local = [event for event in during if event.get("rank") == rank] + nodes = [ + event + for event in local + if event["event"] == "node_evicted" + and event["first_token"] == target_token + and event.get("big_state") is not None + ] + small = [ + event + for event in local + if event["event"] == "small_state_evicted" + and event["first_token"] == target_token + and event.get("small_state") is not None + ] + freed = [event for event in local if event["event"] == "state_freed"] + + def observed_free(event, kind): + return any( + item["pool"] == event[f"{kind}_pool"] + and event[f"{kind}_state"] in item["indexes"] + and item["monotonic"] <= event["monotonic"] + for item in freed + ) + + if not any(observed_free(event, "big") for event in nodes): + raise ValueError(f"No observed target KV/big-state eviction on TP rank {rank}") + if not any(observed_free(event, "small") for event in small): + raise ValueError(f"No observed target small-state eviction on TP rank {rank}") + proof.append(dict(rank=rank, target_nodes=len(nodes), target_small_states=len(small))) + return proof + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--cache-events", type=Path, required=True) + parser.add_argument("--reference", type=Path) + parser.add_argument("--timeout", type=float, default=600) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error("Run through exp") + verified_checkpoint(args.model_directory) + service = json.loads(args.service_metadata.read_text()) + if service.get("framework") != "lightllm" or any( + not service.get(key) for key in ("version", "command", "hardware", "configuration") + ): + raise ValueError("Record complete actual native service provenance") + cases = eviction_cases(args.model_directory) + reference = json.loads(args.reference.read_text()) if args.reference else None + if reference and ( + reference.get("status") not in ("recorded", "passed") + or reference["revision"] != REVISION + or reference["cases"] != cases + ): + raise ValueError("Require a passing reference with identical checkpoint and eviction inputs") + run = Path(os.environ["EXP_RUN_DIR"]) + report = dict( + status="running", + scope=__doc__, + revision=REVISION, + service=service, + cases=cases, + reference=str(args.reference) if args.reference else None, + cache_events=str(args.cache_events), + requests={}, + ) + path = run / "prefix-eviction.json" + + def request(label, case): + result = generate(args.base_url, case, [], args.timeout, run / f"{label}.json") + if reference and result["token_ids"] != reference["requests"][label]["token_ids"]: + raise ValueError(f"Output differs from ordinary reference for {label}") + return result + + def save(label, result): + report["requests"][label] = result + path.write_text(json.dumps(report, indent=2)) + + try: + target = cases["target"] + save("initial", request("initial", target)) + save("warm", request("warm", target)) + initial = report["requests"]["initial"] + warm = report["requests"]["warm"] + if cache_hit(initial) != 0 or cache_hit(warm) < 2048 or initial["token_ids"] != warm["token_ids"]: + raise ValueError("Require a fresh target and exact warm-prefix reuse before pressure") + for name, case in cases.items(): + if name.startswith("small_pressure_"): + save(name, request(name, case)) + save("after_small_pressure", request("after_small_pressure", target)) + small = report["requests"]["after_small_pressure"] + if not 0 < cache_hit(small) < cache_hit(warm) or small["token_ids"] != initial["token_ids"]: + raise ValueError("Small-state pressure must fall back to a shorter big-state prefix with exact output") + for name, case in cases.items(): + if name.startswith("pressure_"): + save(name, request(name, case)) + save("after_pressure", request("after_pressure", target)) + after = report["requests"]["after_pressure"] + if cache_hit(after) >= cache_hit(warm) or after["token_ids"] != initial["token_ids"]: + raise ValueError("Target must lose its warm prefix and preserve exact output after pressure") + save("rewarm", request("rewarm", target)) + if cache_hit(report["requests"]["rewarm"]) < 2048: + raise ValueError("Reallocated target prefix did not become reusable") + with ThreadPoolExecutor(max_workers=4) as executor: + pending = { + f"concurrent_{index}": executor.submit(request, f"concurrent_{index}", target) for index in range(4) + } + for label, future in pending.items(): + save(label, future.result()) + if any(report["requests"][label]["token_ids"] != initial["token_ids"] for label in ["rewarm", *pending]): + raise ValueError("Reallocated target outputs differ under reuse/concurrency") + events = [json.loads(line) for line in args.cache_events.read_text().splitlines() if line] + report["eviction_proof"] = validate_eviction( + events, service["configuration"]["tp"], target["inputs"][0], warm["ended"], after["started"] + ) + report["status"] = "passed" if reference else "recorded" + except Exception as error: + report.update(status="failed", error_type=type(error).__name__, error=str(error)) + raise + finally: + path.write_text(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_retrieval.py b/test/benchmark/qwen38/check_retrieval.py new file mode 100644 index 0000000000..b9ef80fa9d --- /dev/null +++ b/test/benchmark/qwen38/check_retrieval.py @@ -0,0 +1,114 @@ +import argparse +import hashlib +import json +import os +from pathlib import Path +import secrets +import shutil +import time + +import requests + +from lightllm.server.tokenizer import get_tokenizer +from test.benchmark.qwen38.check_generation import verified_checkpoint +from test.benchmark.qwen38.retrieval_quality import load_workload, score_records +from test.benchmark.service.completion import completion_request, completion_tokens + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--framework", choices=("lightllm", "vllm", "sglang"), required=True) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--manifest", type=Path, required=True) + parser.add_argument("--timeout", type=float, default=1200) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + run = Path(os.environ["EXP_RUN_DIR"]) + shutil.copyfile(__file__, run / "client.py") + verified_checkpoint(args.model_directory) + manifest = load_workload(args.manifest, args.model_directory) + service = json.loads(args.service_metadata.read_text()) + if service.get("framework") != args.framework or any( + not service.get(key) for key in ("version", "command", "hardware", "configuration") + ): + raise ValueError("Record actual service framework, version, command, hardware and configuration") + configuration = service["configuration"] + context = configuration.get("context_limit", configuration.get("effective_token_budget", 0)) + if type(context) is not int or context < manifest["context_limit"]: + raise ValueError("Service must declare the actual full context budget") + tokenizer = get_tokenizer(str(args.model_directory)).tokenizer + config = json.loads((args.model_directory / "generation_config.json").read_text()) + eos = config["eos_token_id"] + eos_ids = sorted(set(eos if isinstance(eos, list) else [eos]) | {tokenizer.eos_token_id}) + report = dict( + status="running", + scope=__doc__, + service=service, + manifest_sha256=hashlib.sha256(json.dumps(manifest, sort_keys=True).encode()).hexdigest(), + results={}, + ) + output = run / "long-retrieval.json" + client = requests.Session() + client.trust_env = False + active_case = None + try: + for case in manifest["cases"]: + active_case = case["name"] + endpoint, payload = completion_request( + args.framework, str(args.model_directory.resolve()), case["prompt_ids"], case["max_output_tokens"] + ) + if args.framework == "lightllm": + payload["parameters"]["disable_prompt_cache"] = True + else: + payload["cache_salt"] = secrets.token_urlsafe(32) + started = time.monotonic() + response = client.post(args.base_url.rstrip("/") + endpoint, json=payload, timeout=args.timeout) + ended = time.monotonic() + (run / f"{active_case}.json").write_text(response.text) + response.raise_for_status() + body = response.json() + result = completion_tokens(body, args.framework, case["prompt_ids"], case["max_output_tokens"], eos_ids) + if args.framework == "lightllm": + cached = body["tokens"][0]["prompt_cache_len"] + else: + cached = (body["usage"].get("prompt_tokens_details") or {}).get("cached_tokens") + if type(cached) is not int or cached != 0: + raise ValueError("Every context case must report exactly zero cached prompt tokens") + text = tokenizer.decode(result["token_ids"], skip_special_tokens=True) + score = score_records(text, case["expected_records"]) + score["correct"] = score["correct"] and not result["truncated"] + result.update( + text=text, + prompt_tokens=len(case["prompt_ids"]), + cached_input_tokens=cached, + observed_request_s=ended - started, + **score, + ) + report["results"][active_case] = result + output.write_text(json.dumps(report, indent=2)) + print( + json.dumps( + dict( + case=active_case, + **{key: result[key] for key in ("prompt_tokens", "correct", "records", "truncated")}, + ) + ), + flush=True, + ) + report.update(cases=len(manifest["cases"]), correct=sum(row["correct"] for row in report["results"].values())) + if report["cases"] != 4 or report["correct"] != report["cases"]: + raise ValueError("Some declared-context retrieval answers differ from the pinned known records") + report["status"] = "passed" + except Exception as error: + report.update(status="failed", failed_case=active_case, error_type=type(error).__name__, error=str(error)) + raise + finally: + output.write_text(json.dumps(report, indent=2)) + client.close() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_scheduler.py b/test/benchmark/qwen38/check_scheduler.py new file mode 100644 index 0000000000..37b17b129d --- /dev/null +++ b/test/benchmark/qwen38/check_scheduler.py @@ -0,0 +1,215 @@ +import argparse +from concurrent.futures import ThreadPoolExecutor, as_completed +import hashlib +import json +import os +from pathlib import Path +import threading +import time + +import requests +from transformers import AutoTokenizer + +from check_generation import REVISION, TOKENIZER_SHA256, cases_for, verified_checkpoint + + +def scheduler_cases(root): + raw = cases_for(root) + tokenizer = AutoTokenizer.from_pretrained(root, local_files_only=True) + result = { + name: {"inputs": raw[name], "max_new_tokens": length, "ignore_eos": True} + for name, length in ( + ("english", 32), + ("chinese", 31), + ("code_unicode", 17), + ("long_2049", 32), + ("long_4063", 32), + ) + } + result["eos_history"] = { + "inputs": [248044] + raw["english"] + [248046] + raw["chinese"], + "max_new_tokens": 19, + "ignore_eos": True, + } + for name, question in ( + ("chat_ok", "Reply with exactly OK."), + ("chat_number", "What is 2+2? Reply only with the number."), + ): + result[name] = { + "inputs": tokenizer.apply_chat_template( + [{"role": "user", "content": question}], + tokenize=True, + add_generation_prompt=True, + enable_thinking=False, + ), + "max_new_tokens": 96, + "ignore_eos": False, + } + generation_eos = json.loads((root / "generation_config.json").read_text())["eos_token_id"] + eos = set(generation_eos if isinstance(generation_eos, list) else [generation_eos]) + eos.add(tokenizer.eos_token_id) + return result, sorted(eos) + + +def generate(base_url, case, eos_ids, timeout, response_path, cold=False): + parameters = {key: value for key, value in case.items() if key != "inputs"} + parameters.update(do_sample=False, return_details=True, print_eos_token=True, disable_prompt_cache=cold) + started = time.monotonic() + with requests.Session() as client: + client.trust_env = False + response = client.post( + base_url.rstrip("/") + "/generate", + json={"inputs": case["inputs"], "parameters": parameters}, + timeout=timeout, + ) + ended = time.monotonic() + response_path.write_text(response.text) + response.raise_for_status() + body = response.json() + ids = validate_generation(body, len(case["inputs"]), case["max_new_tokens"], case["ignore_eos"], eos_ids) + return {"token_ids": ids, "response": body, "started": started, "ended": ended} + + +def validate_generation(body, prompt_tokens, limit, ignore_eos, eos_ids): + ids = [token["id"] for token in body["tokens"]] + if not ids or any(type(token) is not int or token < 0 for token in ids) or len(ids) > limit: + raise ValueError("Invalid output token IDs or output length") + if body["prompt_tokens"] != prompt_tokens or body["count_output_tokens"] != len(ids): + raise ValueError("Service prompt/output counts do not match exact token IDs") + if ignore_eos: + if len(ids) != limit or body["finish_reason"] != "length": + raise ValueError("Fixed-length request did not reach its requested length") + elif ids[-1] not in eos_ids or any(token in eos_ids for token in ids[:-1]) or body["finish_reason"] != "stop": + raise ValueError("EOS request must stop at its first declared EOS token") + return ids + + +def monitor(base_url, stop, summary, path): + with requests.Session() as client: + client.trust_env = False + while not stop.is_set(): + try: + response = client.get(base_url.rstrip("/") + "/metrics", timeout=5) + response.raise_for_status() + path.write_text(response.text) + for line in response.text.splitlines(): + if line.startswith(("lightllm_batch_current_size", "lightllm_num_running_reqs")): + summary["maximum_router_batch"] = max(summary["maximum_router_batch"], float(line.split()[-1])) + summary["samples"] += 1 + except Exception as error: + summary["last_error"] = str(error) + stop.wait(0.02) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--reference", type=Path) + parser.add_argument("--concurrency", type=int, nargs="+", default=[2, 4, 8]) + parser.add_argument("--timeout", type=float, default=600) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + if not args.concurrency or min(args.concurrency) < 2 or max(args.concurrency) > 8: + parser.error("Concurrent waves must use between two and eight requests") + run = Path(os.environ["EXP_RUN_DIR"]) + verified_checkpoint(args.model_directory) + metadata = json.loads(args.service_metadata.read_text()) + if metadata.get("framework") != "lightllm" or any( + not metadata.get(key) for key in ("version", "command", "hardware", "configuration") + ): + raise ValueError("Record the actual native service version, command, hardware and configuration") + cases, eos_ids = scheduler_cases(args.model_directory) + source = { + name: hashlib.sha256((args.model_directory / name).read_bytes()).hexdigest() + for name in ("tokenizer.json", "tokenizer_config.json", "generation_config.json") + } + reference = json.loads(args.reference.read_text()) if args.reference else None + if reference and ( + reference.get("status") not in ("passed", "recorded") + or reference["revision"] != REVISION + or reference["cases"] != cases + or reference["source"] != source + or reference["eos_ids"] != eos_ids + ): + raise ValueError("Reference must pass with identical checkpoint, tokenizer, inputs and EOS contract") + report = dict( + status="running", + scope="full trained scheduler/EOS/cache correctness; no performance claim", + revision=REVISION, + tokenizer_sha256=TOKENIZER_SHA256, + source=source, + service=metadata, + cases=cases, + eos_ids=eos_ids, + reference=str(args.reference) if args.reference else None, + serial={}, + waves=[], + mismatches=[], + ) + output = run / "scheduler.json" + stop = threading.Event() + metrics = {"maximum_router_batch": 0, "samples": 0} + thread = threading.Thread(target=monitor, args=(args.base_url, stop, metrics, run / "metrics.txt"), daemon=True) + thread.start() + active = None + try: + for name, case in cases.items(): + active = {"wave": "serial", "case": name} + result = generate(args.base_url, case, eos_ids, args.timeout, run / f"serial-{name}.txt") + report["serial"][name] = result + if reference and result["token_ids"] != reference["serial"][name]["token_ids"]: + report["mismatches"].append({"wave": "serial_reference", "case": name}) + output.write_text(json.dumps(report, indent=2)) + names = list(cases) + plans = [(count, False) for count in args.concurrency] + [(8, True), (8, False)] + for index, (count, cold) in enumerate(plans): + order = names[index:] + names[:index] + if index == len(plans) - 1: + order = list(reversed(order)) + wave = {"concurrency": count, "disable_prompt_cache": cold, "results": {}} + report["waves"].append(wave) + with ThreadPoolExecutor(max_workers=count) as executor: + pending = { + executor.submit( + generate, + args.base_url, + cases[name], + eos_ids, + args.timeout, + run / f"wave-{index}-{name}.txt", + cold, + ): name + for name in order + } + for future in as_completed(pending): + name = pending[future] + active = {"wave": index, "case": name} + result = future.result() + wave["results"][name] = result + if result["token_ids"] != report["serial"][name]["token_ids"]: + report["mismatches"].append({"wave": index, "case": name}) + output.write_text(json.dumps(report, indent=2)) + print( + json.dumps({"wave": index, "concurrency": count, "cold": cold, "mismatches": report["mismatches"]}), + flush=True, + ) + if report["mismatches"]: + raise ValueError("Serial/concurrent or ordinary/speculative token IDs differ") + if metrics["maximum_router_batch"] < 2: + raise ValueError("No actual multi-request router batch observed; packed serving remains unverified") + report["status"] = "passed" if reference else "recorded" + except Exception as error: + report.update(status="failed", failed_request=active, error_type=type(error).__name__, error=str(error)) + raise + finally: + stop.set() + thread.join(timeout=6) + report["router_metrics"] = metrics + output.write_text(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_streaming.py b/test/benchmark/qwen38/check_streaming.py new file mode 100644 index 0000000000..196d75edfb --- /dev/null +++ b/test/benchmark/qwen38/check_streaming.py @@ -0,0 +1,76 @@ +import argparse +import hashlib +import json +import os +from pathlib import Path + +from check_generation import REVISION, verified_checkpoint +from test.benchmark.service.token_stream import request_stream + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--framework", choices=("lightllm", "vllm", "sglang"), required=True) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--cases", nargs="+", default=["english", "chinese", "long_2049"]) + parser.add_argument("--timeout", type=float, default=600) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error("Run through exp") + verified_checkpoint(args.model_directory) + reference = json.loads(args.reference.read_text()) + service = json.loads(args.service_metadata.read_text()) + if service.get("framework") != args.framework or any( + not service.get(key) for key in ("version", "command", "hardware", "configuration") + ): + raise ValueError("Record actual service provenance") + if reference["status"] not in ("passed", "recorded") or reference["revision"] != REVISION: + raise ValueError("Require a passing pinned scheduler reference") + for name, expected in reference["source"].items(): + if hashlib.sha256((args.model_directory / name).read_bytes()).hexdigest() != expected: + raise ValueError(f"Checkpoint/tokenizer changed after reference: {name}") + run = Path(os.environ["EXP_RUN_DIR"]) + report = dict( + status="running", + scope="real streaming protocol/greedy parity; no performance comparison", + revision=REVISION, + service=service, + reference=str(args.reference), + results={}, + ) + path = run / "streaming.json" + active = None + try: + for name in args.cases: + active = name + case = reference["cases"][name] + if not case["ignore_eos"]: + raise ValueError("Streaming fixed-length gate requires EOS ignored") + result = request_stream( + args.base_url, + args.framework, + str(args.model_directory.resolve()), + case["inputs"], + case["max_new_tokens"], + args.timeout, + run / f"response-{name}.sse", + ) + result["reference_matches"] = result["token_ids"] == reference["serial"][name]["token_ids"] + report["results"][name] = result + path.write_text(json.dumps(report, indent=2)) + print(json.dumps({"case": name, **result}), flush=True) + if not all(result["reference_matches"] for result in report["results"].values()): + raise ValueError("Streamed token IDs differ from the exact scheduler reference") + report["status"] = "passed" + except Exception as error: + report.update(status="failed", failed_case=active, error_type=type(error).__name__, error=str(error)) + raise + finally: + path.write_text(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/generation_observations.py b/test/benchmark/qwen38/generation_observations.py new file mode 100644 index 0000000000..cebf1130b3 --- /dev/null +++ b/test/benchmark/qwen38/generation_observations.py @@ -0,0 +1,48 @@ +from test.benchmark.qwen38.check_generation import REVISION +from test.benchmark.service.stream_cohort import StreamingWorkload + + +def workloads_from_observations(report, expected_cases): + if ( + report.get("status") != "observations_recorded" + or report.get("service", {}).get("framework") not in ("vllm", "sglang") + or report.get("revision") != REVISION + or report.get("cases") != expected_cases + or report.get("reference") is not None + or report.get("cache_policy") != "cold" + ): + raise ValueError("Require explicit cold baseline observations with pinned inputs") + results, length = report.get("results", {}), report.get("output_length") + if type(length) is not int or length < 1 or set(results) != set(expected_cases): + raise ValueError("Require complete baseline observations and a positive output length") + workloads, differences = [], [] + for name, prompt in expected_cases.items(): + result = results[name] + ids, body = result["token_ids"], result["response"] + usage, choices = body["usage"], body["choices"] + cached = (usage.get("prompt_tokens_details") or {}).get("cached_tokens") + if ( + len(choices) != 1 + or choices[0]["token_ids"] != ids + or usage["prompt_tokens"] != len(prompt) + or usage["completion_tokens"] != length + or len(ids) != length + or type(cached) is not int + or cached != 0 + ): + raise ValueError("Baseline observations must preserve raw IDs, counts and explicit zero cache") + work = StreamingWorkload(name, prompt, ids) + if name.startswith("repeat_"): + original = name.removeprefix("repeat_") + if original not in expected_cases or prompt != expected_cases[original]: + raise ValueError("Repeat observations require identical pinned inputs") + if ids != results[original]["token_ids"]: + differences.append(int(original.removeprefix("long_"))) + else: + workloads.append(work) + quality = dict(status="failed" if differences else "passed", mismatches=differences) + if quality != report.get("repeat_consistency"): + raise ValueError("Baseline repeat-consistency summary differs from actual observations") + if not workloads: + raise ValueError("Require a non-repeated workload") + return workloads, quality diff --git a/test/benchmark/qwen38/math_quality.py b/test/benchmark/qwen38/math_quality.py new file mode 100644 index 0000000000..fd7bbc8c8d --- /dev/null +++ b/test/benchmark/qwen38/math_quality.py @@ -0,0 +1,127 @@ +from dataclasses import dataclass +from decimal import Decimal +import hashlib +import json +import random +import re + +DATASET_REVISION = "3101c7d5072418e28b9008a6636bde82a006892c" +DATASET_SHA256 = "3730d312f6e3440559ace48831e51066acaca737f6eabec99bccb9e4b3c39d14" +NUMBER = re.compile(r"[+-]?(?:\d{1,3}(?:,\d{3})+|\d+)(?:\.\d+)?") +INSTRUCTION = ( + "Solve the following math problem. Show your reasoning briefly, then end with a final line " + "exactly in the format #### . Use a plain decimal number.\n\nProblem:\n" +) + + +def normalize_number(value): + if not isinstance(value, str) or not NUMBER.fullmatch(value.strip()): + return None + number = Decimal(value.strip().replace(",", "")) + if number == 0: + return "0" + text = format(number, "f") + return text.rstrip("0").rstrip(".") if "." in text else text + + +def extract_final_number(completion): + if not isinstance(completion, str): + return None + lines = completion.rstrip().splitlines() + if not lines or sum(line.strip().startswith("####") for line in lines) != 1: + return None + match = re.fullmatch(r"####[ \t]+(.+)", lines[-1].strip()) + return normalize_number(match[1]) if match else None + + +def score_answer(completion, expected): + gold = normalize_number(expected) + if gold is None: + raise ValueError("Expected answer must be a finite plain decimal number") + predicted = extract_final_number(completion) + return dict(expected=gold, predicted=predicted, correct=predicted is not None and predicted == gold) + + +def compare_outputs(reference, actual): + digest = reference.get("manifest_sha256") + if ( + reference.get("status") != "recorded" + or actual.get("status") != "recorded" + or not isinstance(digest, str) + or not re.fullmatch(r"[0-9a-f]{64}", digest) + or actual.get("manifest_sha256") != digest + ): + raise ValueError("Require complete scored runs with identical workload manifests") + + def rows(report): + values = report["results"] + if type(report.get("scored")) is not int or report["scored"] < 1 or len(values) != report["scored"]: + raise ValueError("Scored count differs from completed cases") + indices = [row["dataset_index"] for row in values] + if any(type(index) is not int or index < 0 for index in indices) or len(set(indices)) != len(indices): + raise ValueError("Require distinct actual dataset indices") + return {row["dataset_index"]: row for row in values} + + expected, observed = rows(reference), rows(actual) + if set(expected) != set(observed): + raise ValueError("Scored runs contain different selected cases") + fields = ("token_ids", "finish_reason", "truncated", "correct", "predicted") + mismatches = [ + index for index in sorted(expected) if any(expected[index][key] != observed[index][key] for key in fields) + ] + return dict(status="passed" if not mismatches else "failed", cases=len(expected), mismatches=mismatches) + + +@dataclass(frozen=True) +class MathExample: + index: int + question: str + expected: str + + +def load_examples(path): + payload = path.read_bytes() + if hashlib.sha256(payload).hexdigest() != DATASET_SHA256: + raise ValueError("GSM8K test data differs from the pinned source") + examples = [] + for index, line in enumerate(payload.splitlines()): + row = json.loads(line) + question = row.get("question") + answer = extract_final_number(row.get("answer")) + if not isinstance(question, str) or not question.strip() or answer is None: + raise ValueError("Dataset requires a question and unambiguous final numeric answer") + examples.append(MathExample(index, question, answer)) + if not examples or len({example.question for example in examples}) != len(examples): + raise ValueError("Dataset requires distinct questions") + return tuple(examples) + + +def select_examples(examples, limit, seed): + if type(limit) is not int or not 1 <= limit <= len(examples) or type(seed) is not int: + raise ValueError("Require a valid sample count and integer seed") + return tuple(examples[index] for index in random.Random(seed).sample(range(len(examples)), limit)) + + +def prepare_cases(examples, tokenizer, max_output_tokens=512, context_limit=4096): + if any(type(value) is not int or value < 1 for value in (max_output_tokens, context_limit)): + raise ValueError("Output and context limits must be positive integers") + cases = [] + for example in examples: + messages = [{"role": "user", "content": INSTRUCTION + example.question}] + ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, enable_thinking=False) + if not isinstance(ids, list) or not ids or any(type(token) is not int or token < 0 for token in ids): + raise ValueError("Tokenizer must return actual nonnegative integer prompt IDs") + if len(ids) + max_output_tokens > context_limit: + raise ValueError("Math example exceeds the declared serving context") + cases.append( + dict( + dataset_index=example.index, + question_sha256=hashlib.sha256(example.question.encode()).hexdigest(), + prompt_ids=ids, + expected_number=example.expected, + max_output_tokens=max_output_tokens, + ) + ) + if not cases or len({case["dataset_index"] for case in cases}) != len(cases): + raise ValueError("Require distinct selected examples") + return cases diff --git a/test/benchmark/qwen38/measure_aiperf.py b/test/benchmark/qwen38/measure_aiperf.py new file mode 100644 index 0000000000..fc61f16102 --- /dev/null +++ b/test/benchmark/qwen38/measure_aiperf.py @@ -0,0 +1,310 @@ +import argparse +import hashlib +import json +import math +import os +from pathlib import Path +import shutil +import statistics +import subprocess + +import requests + + +def tokenizer_cache(model, directory, revision): + snapshot = directory / "models--Qwen--Qwen3.8-Flash-Next" / "snapshots" / revision + snapshot.mkdir(parents=True, exist_ok=True) + hashes = {} + for name in ( + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "added_tokens.json", + ): + original = model.resolve() / name + if not original.exists(): + if name in ("config.json", "tokenizer.json", "tokenizer_config.json"): + raise ValueError(f"Missing pinned tokenizer file: {name}") + continue + destination = snapshot / name + if not destination.exists(): + destination.symlink_to(original) + if destination.read_bytes() != original.read_bytes(): + raise ValueError("Cached tokenizer differs from the pinned checkpoint") + hashes[name] = hashlib.sha256(original.read_bytes()).hexdigest() + return hashes + + +def profiling_records(directory): + records = [json.loads(line) for line in (directory / "profile_export_raw.jsonl").read_text().splitlines()] + records = [row for row in records if row["metadata"]["benchmark_phase"] == "profiling"] + + def index(row): + identity = row["metadata"]["conversation_id"] + if not identity.startswith("session_") or not identity.removeprefix("session_").isdigit(): + raise ValueError("Unexpected source conversation identity") + return int(identity.removeprefix("session_")) + + records.sort(key=index) + if [index(row) for row in records] != list(range(len(records))): + raise ValueError("Missing or duplicated source conversations") + return records + + +def audit_records(directory, trace): + summary = json.loads((directory / "profile_export_aiperf.json").read_text()) + records = profiling_records(directory) + if len(records) != len(trace): + raise ValueError("Incomplete replay") + if sorted(row["metadata"]["session_num"] for row in records) != list(range(len(trace))): + raise ValueError("Missing, duplicated or reordered dataset sessions") + requests_for_comparison, starts, cache_counts, empty_content = [], [], [], [] + for index, (record, expected) in enumerate(zip(records, trace)): + payload, metadata = record["payload"], record["metadata"] + if record["status"] != 200 or metadata.get("was_cancelled"): + raise ValueError("Failed or cancelled request") + if payload["max_tokens"] != expected["output_length"] or not payload.get("ignore_eos"): + raise ValueError("Replay changed the specified output work") + usage, done, finished, has_text = None, False, False, False + for response in record["responses"]: + for packet in response.get("packets", []): + if packet.get("name") != "data": + continue + if packet["value"] == "[DONE]": + done = True + continue + data = json.loads(packet["value"]) + usage = data.get("usage") or usage + finished |= any(choice.get("finish_reason") == "length" for choice in data.get("choices", [])) + has_text |= any(bool(choice.get("text")) for choice in data.get("choices", [])) + if not done or not finished or not usage: + raise ValueError("Missing stream terminator, length finish or actual usage") + if type(usage.get("completion_tokens")) is not int or usage["completion_tokens"] != expected["output_length"]: + raise ValueError("Actual completion count differs from the trace") + if ( + type(usage.get("prompt_tokens")) is not int + or not 0 < usage["prompt_tokens"] + usage["completion_tokens"] <= 4096 + ): + raise ValueError("Actual request exceeds the shared context limit") + cache_counts.append((usage.get("prompt_tokens_details") or {}).get("cached_tokens")) + if not has_text: + empty_content.append(index) + requests_for_comparison.append( + dict( + prompt=payload["prompt"], + max_tokens=payload["max_tokens"], + temperature=payload["temperature"], + ignore_eos=payload["ignore_eos"], + ) + ) + starts.append(metadata["request_start_ns"]) + + errors = summary.get("error_summary", []) + for error in errors: + detail = error["error_details"] + if ( + detail["type"] != "InvalidInferenceResultError" + or "No responses with actual content" not in detail["message"] + ): + raise ValueError("AIPerf reported a transport, protocol or other request error") + if sum(error["count"] for error in errors) != len(empty_content): + raise ValueError("AIPerf errors do not match the complete empty-text responses") + if summary["request_count"]["avg"] != len(trace) - len(empty_content): + raise ValueError("AIPerf successful request count mismatch") + successful_tokens = sum(row["output_length"] for index, row in enumerate(trace) if index not in empty_content) + if summary["output_token_count"]["sum"] != successful_tokens: + raise ValueError("AIPerf successful token count mismatch") + + lag = [(start - starts[0]) / 1e6 - (row["timestamp"] - trace[0]["timestamp"]) for start, row in zip(starts, trace)] + required = ( + "request_throughput", + "output_token_throughput", + "time_to_first_token", + "inter_token_latency", + "request_latency", + ) + for name in required: + if not math.isfinite(summary[name]["avg"]): + raise ValueError(f"Invalid AIPerf metric: {name}") + return dict( + requests=len(records), + output_tokens=sum(row["output_length"] for row in trace), + content_errors=len(empty_content), + error_rate=len(empty_content) / len(trace), + empty_content_sessions=empty_content, + successful_output_tokens=successful_tokens, + metric_scope="AIPerf successful content-bearing responses; empty-text responses remain errors", + request_manifest_sha256=hashlib.sha256( + json.dumps(requests_for_comparison, sort_keys=True).encode() + ).hexdigest(), + cached_tokens=cache_counts, + relative_schedule_lag_ms=dict(median=statistics.median(lag), max=max(lag), min=min(lag)), + metrics={name: summary[name] for name in required}, + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--trace-directory", type=Path, required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--aiperf", type=Path, required=True) + parser.add_argument("--speed", type=float, default=1.0) + parser.add_argument( + "--fresh-owned-service", + action="store_true", + help="Native replay in a newly owned service with no prior workload requests", + ) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + if not math.isfinite(args.speed) or args.speed <= 0: + parser.error("Speed must be finite and positive") + service = json.loads(args.service_metadata.read_text()) + for key in ( + "framework", + "version", + "software", + "command", + "hardware", + "configuration", + "base_url", + "model_revision", + ): + if not service.get(key): + raise ValueError(f"Missing actual service provenance: {key}") + if service["status"] not in ("ready", "measuring_aiperf"): + raise ValueError("AIPerf requires an exclusively owned, ready service") + manifest = json.loads((args.trace_directory / "manifest.json").read_text()) + raw = (args.trace_directory / "replay.jsonl").read_bytes() + if hashlib.sha256(raw).hexdigest() != manifest["trace_sha256"]: + raise ValueError("Trace checksum mismatch") + trace = [json.loads(line) for line in raw.splitlines()] + if len(trace) != manifest["requests"] or manifest["context_length"] != 4096: + raise ValueError("Trace manifest mismatch") + run = Path(os.environ["EXP_RUN_DIR"]) + shutil.copyfile(__file__, run / Path(__file__).name) + version = subprocess.check_output([str(args.aiperf), "--version"], text=True).strip() + if version != "0.9.0": + raise ValueError("Comparison is pinned to AIPerf 0.9.0") + for row in trace: + row["timestamp"] /= args.speed + trace_path = run / "trace.jsonl" + trace_path.write_text("".join(json.dumps(row) + "\n" for row in trace)) + report = dict( + status="running", + service=service, + aiperf_version=version, + manifest=manifest, + speed=args.speed, + concurrency_ceiling=8, + cache_policy="flush before each replay; prefix reuse within trace enabled", + precision_scope="Performance only; reuse independent model task-quality results", + comparison_eligible=False, + ) + + def save(): + (run / "aiperf-replay.json").write_text(json.dumps(report, indent=2) + "\n") + + try: + save() + cache = args.trace_directory.resolve() / "tokenizer-cache" + report["tokenizer_sha256"] = tokenizer_cache(args.model_directory, cache, service["model_revision"]) + environment = dict( + os.environ, + HF_HUB_CACHE=str(cache), + HUGGINGFACE_HUB_CACHE=str(cache), + HF_HUB_OFFLINE="1", + TRANSFORMERS_OFFLINE="1", + ) + framework = service["framework"] + if framework not in ("lightllm", "vllm", "sglang"): + raise ValueError("Unsupported framework") + reset = "/reset_prefix_cache" if framework == "vllm" else "/flush_cache" + if args.fresh_owned_service: + import psutil + + if framework != "lightllm" or not service["configuration"].get("fresh_aiperf_replay"): + raise ValueError("Fresh-instance cache isolation is restricted to the native replay owner") + if not service.get("fresh_service_no_quality_requests") or service.get("checks"): + raise ValueError("Fresh replay requires a service with no preceding workload requests") + owner = psutil.Process(service["launcher_pid"]) + allowed = {"qwen38-h100-native-fresh-service.py", "qwen38-h100-aiperf-profile.py"} + if owner.create_time() != service["launcher_created"] or not any( + Path(value).name in allowed for value in owner.cmdline() + ): + raise ValueError("Fresh replay owner identity mismatch") + report[ + "cache_policy" + ] = "New owned native instance per replay; no prior workload requests; prefix reuse within trace enabled" + report["cache_reset"] = dict( + method="fresh_owned_service", container=service["container"]["name"], owner_pid=owner.pid + ) + save() + else: + with requests.Session() as session: + session.trust_env = False + response = session.post(service["base_url"] + reset, timeout=120) + report["cache_reset"] = dict(status=response.status_code, body=response.text) + save() + response.raise_for_status() + if response.text.strip().lower() == "false": + raise ValueError("Server refused cache reset") + if framework == "vllm" and response.json().get("success") is not True: + raise ValueError("vLLM did not confirm prefix-cache reset") + command = [ + str(args.aiperf), + "profile", + "--model", + str(args.model_directory.resolve()), + "--tokenizer", + "Qwen/Qwen3.8-Flash-Next", + "--tokenizer-revision", + service["model_revision"], + "--url", + service["base_url"], + "--endpoint-type", + "completions", + "--streaming", + "--input-file", + str(trace_path), + "--custom-dataset-type", + "mooncake_trace", + "--fixed-schedule", + "--fixed-schedule-auto-offset", + "--concurrency", + "8", + "--request-count", + str(len(trace)), + "--random-seed", + "42", + "--use-server-token-count", + "--export-level", + "raw", + "--artifact-dir", + str(run / "aiperf"), + "--ui", + "none", + "--workers-max", + "4", + "--no-auto-plot", + "--request-timeout-seconds", + "600", + "--extra-inputs", + json.dumps(dict(temperature=0, ignore_eos=True)), + ] + report["command"] = command + save() + subprocess.run(command, check=True, env=environment) + report["audit"] = audit_records(run / "aiperf", trace) + report["status"] = "aiperf_replay_recorded" + except BaseException as error: + report.update(status="failed", error_type=type(error).__name__, error=str(error)) + raise + finally: + save() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/measure_fixed_work.py b/test/benchmark/qwen38/measure_fixed_work.py new file mode 100644 index 0000000000..c072a790e2 --- /dev/null +++ b/test/benchmark/qwen38/measure_fixed_work.py @@ -0,0 +1,152 @@ +import argparse +import hashlib +import json +import os +from pathlib import Path +import secrets +import shutil + +from test.benchmark.qwen38.check_cold_streams import workloads_from_reference +from test.benchmark.qwen38.check_generation import cases_for, verified_checkpoint +from test.benchmark.qwen38.generation_observations import workloads_from_observations +from test.benchmark.service.stream_cohort import run_cohort, summarize_cohort +from test.benchmark.service.token_stream import request_stream + + +def summarize_fixed_work(cohort, workloads): + expected = {work.name: work for work in workloads} + results = [] + for record in cohort["records"]: + if record["status"] != "passed" and record.get("error") != "Output IDs differ from the qualified reference": + raise ValueError("Incomplete or invalid streaming request") + work = expected[record["workload"]] + result = record["result"] + metrics = result["metrics"] + if metrics["input_tokens"] != len(work.prompt_ids) or len(result["token_ids"]) != len(work.expected_ids): + raise ValueError("Actual input/output lengths differ from the workload") + cached = metrics.get("cached_input_tokens") + if type(cached) is not int or cached != 0: + raise ValueError("Every measured request must explicitly report zero cached input tokens") + results.append(result) + if len(results) != len(workloads): + raise ValueError("Missing workload requests") + return summarize_cohort(results) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--concurrency", type=int, nargs="+", default=[1, 2, 4, 8]) + parser.add_argument( + "--observed-generation", + action="store_true", + help="Measure count-valid baseline observations; retain repeat-quality failures", + ) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + if not 2 <= args.repeats <= 20 or any(c not in (1, 2, 4, 8) for c in args.concurrency): + parser.error("Require 2–20 repeats and concurrency in 1/2/4/8") + service = json.loads(args.service_metadata.read_text()) + for key in ("framework", "version", "software", "command", "hardware", "configuration", "base_url"): + if not service.get(key): + raise ValueError(f"Missing actual service provenance: {key}") + if service["framework"] not in ("lightllm", "vllm", "sglang"): + raise ValueError("Unsupported framework") + verified_checkpoint(args.model_directory) + reference = json.loads(args.reference.read_text()) + lengths = [int(name.removeprefix("long_")) for name in reference["cases"] if name.startswith("long_")] + expected_cases = cases_for(args.model_directory, lengths) + if args.observed_generation: + if ( + service["framework"] not in ("vllm", "sglang") + or reference.get("service", {}).get("framework") != service["framework"] + ): + raise ValueError("Observation timing is restricted to matching baseline frameworks") + workloads, reference_quality = workloads_from_observations(reference, expected_cases) + else: + workloads = workloads_from_reference(reference, expected_cases) + reference_quality = dict(status="qualified_generation") + run = Path(os.environ["EXP_RUN_DIR"]) + shutil.copyfile(__file__, run / Path(__file__).name) + manifest = [dict(name=w.name, prompt_ids=w.prompt_ids, output_length=len(w.expected_ids)) for w in workloads] + report = dict( + status="running", + scope=__doc__, + service=service, + manifest=manifest, + manifest_sha256=hashlib.sha256(json.dumps(manifest, sort_keys=True).encode()).hexdigest(), + reference_sha256=hashlib.sha256(args.reference.read_bytes()).hexdigest(), + reference_kind="observations" if args.observed_generation else "qualified_generation", + reference_quality=reference_quality, + service_sha256=hashlib.sha256(args.service_metadata.read_bytes()).hexdigest(), + repeats=args.repeats, + cohorts=[], + warmups=[], + comparison_eligible=False, + source_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + ) + + def save(): + (run / "fixed-work.json").write_text(json.dumps(report, indent=2)) + + def request(work, path): + options = ( + dict(disable_prompt_cache=True) + if service["framework"] == "lightllm" + else dict(cache_salt=secrets.token_urlsafe(32)) + ) + return request_stream( + service["base_url"], + service["framework"], + str(args.model_directory.resolve()), + list(work.prompt_ids), + len(work.expected_ids), + 600, + path, + **options, + ) + + try: + save() + + for concurrency in args.concurrency: + cohort = run_cohort(workloads * 4, concurrency, request, run / f"warmup-{concurrency}") + report["warmups"].append(cohort) + summarize_fixed_work(cohort, workloads * 4) + save() + for repeat in range(args.repeats): + order = ( + args.concurrency[repeat % len(args.concurrency) :] + args.concurrency[: repeat % len(args.concurrency)] + ) + sequence = workloads[repeat % len(workloads) :] + workloads[: repeat % len(workloads)] + for concurrency in order: + cohort = run_cohort(sequence * 4, concurrency, request, run / f"repeat-{repeat}-c{concurrency}") + cohort["repeat"] = repeat + report["cohorts"].append(cohort) + cohort["fixed_work_metrics"] = summarize_fixed_work(cohort, sequence * 4) + save() + print( + json.dumps( + dict( + repeat=repeat, + concurrency=concurrency, + exact_status=cohort["status"], + output_tokens_per_s=cohort["fixed_work_metrics"]["output_tokens_per_s"], + ) + ), + flush=True, + ) + report["status"] = "fixed_work_measured" + except BaseException as error: + report.update(status="failed", error_type=type(error).__name__, error=str(error)) + raise + finally: + save() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/prepare_agentx.py b/test/benchmark/qwen38/prepare_agentx.py new file mode 100644 index 0000000000..e1c0d0eb0a --- /dev/null +++ b/test/benchmark/qwen38/prepare_agentx.py @@ -0,0 +1,190 @@ +import argparse +import hashlib +import importlib.metadata +import json +import os +from pathlib import Path + + +VERSION = "1.0.13" +DATASET_REVISION = "8fecd2fc56694469f758f0afbbb6335ad3043740" +DATASET_SHA256 = "e39cd2ff3eba21d4a3664be51da743ac3d2149a1933898cafc7bfeac8147eeef" + + +def digest(value): + return hashlib.sha256(value).hexdigest() + + +def canonical(value): + return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode() + + +def prompt_length(tokenizer, messages): + encoded = tokenizer.apply_chat_template(messages, add_generation_prompt=True) + return len(encoded["input_ids"] if hasattr(encoded, "keys") else encoded) + + +def tool_identity(): + import aiperf + import agentx_cli + + version = importlib.metadata.version("agentx-harness") + if version != VERSION: + raise ValueError(f"Require agentx-harness {VERSION}, got {version}") + package = Path(aiperf.__file__).parent + files = {str(p.relative_to(package)): digest(p.read_bytes()) for p in sorted(package.rglob("*.py"))} + files["../agentx_cli.py"] = digest(Path(agentx_cli.__file__).read_bytes()) + dependencies = sorted((d.metadata["Name"], d.version) for d in importlib.metadata.distributions()) + return dict(version=version, source_sha256=digest(canonical(files)), files=files, dependencies=dependencies) + + +def frozen_sessions(conversations, model, streams, turns): + sessions, sources = [], [] + for conversation in sorted(conversations, key=lambda c: c.session_id): + if len(conversation.turns) < turns: + continue + messages, payloads = [], [] + for turn in conversation.turns[:turns]: + if turn.reset_context: + messages = [] + messages.extend(turn.raw_messages) + payloads.append( + dict( + model=model, + messages=json.loads(json.dumps(messages)), + max_tokens=turn.max_tokens, + temperature=0, + seed=42, + ignore_eos=True, + stream=True, + stream_options=dict(include_usage=True), + ) + ) + sessions.append(dict(session_id=conversation.session_id, payloads=payloads)) + sources.append( + dict( + session_id=conversation.session_id, + turns=[ + dict( + source_trace_id=t.source_trace_id, + source_outer_idx=t.source_outer_idx, + source_kind=t.source_kind, + output_tokens=t.max_tokens, + payload_sha256=digest(canonical(p)), + ) + for t, p in zip(conversation.turns[:turns], payloads) + ], + ) + ) + if len(sessions) == streams: + return dict(data=sessions), sources + raise ValueError(f"Only {len(sessions)} eligible streams, require {streams}") + + +def main(): + parser = argparse.ArgumentParser(description="Freeze AgentX and a fixed Weka control workload") + parser.add_argument("--source", type=Path, required=True, help="Pinned HF snapshot traces.jsonl") + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--model", default="Qwen3.8-Flash-Next") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--streams", type=int, default=32) + parser.add_argument("--turns", type=int, default=8) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + if not 1 <= args.streams <= 256 or not 1 <= args.turns <= 100: + parser.error("Invalid stream or turn count") + if DATASET_REVISION not in args.source.parts: + parser.error("Use the pinned HF snapshot, not refs/main") + + from aiperf.common import random_generator as rng + from aiperf.common.tokenizer import Tokenizer + from aiperf.config import BenchmarkConfig, BenchmarkRun + from aiperf.dataset.generator.prompt import PromptGenerator + from aiperf.dataset.loader.weka_trace import WekaTraceLoader, _trace_peak_context_length + from aiperf.dataset.loader.weka_trace_models import WekaTrace + + identity = tool_identity() + raw = args.source.read_bytes() + if digest(raw) != DATASET_SHA256: + raise ValueError("Dataset bytes differ from the pinned snapshot") + traces = [WekaTrace.model_validate(json.loads(line)) for line in raw.splitlines() if line.strip()] + eligible = [t for t in traces if _trace_peak_context_length(t, max_osl=None) <= 262144] + cfg = BenchmarkConfig.model_validate( + dict( + models=[args.model], + endpoint=dict(urls=["http://127.0.0.1:28081"], type="chat", streaming=True), + datasets=[dict(name="main", type="file", path=str(args.source), format="weka_trace")], + tokenizer=dict(name=str(args.model_directory)), + phases=[dict(name="profiling", type="concurrency", concurrency=32, requests=1)], + ) + ) + rng.init(42) + tokenizer = Tokenizer.from_pretrained(str(args.model_directory), resolve_alias=False) + generator = PromptGenerator(prompts=None, prefix_prompts=None, tokenizer=tokenizer) + run = BenchmarkRun(benchmark_id="qwen38-fixed-work-v1", cfg=cfg, random_seed=42, artifact_dir=args.output) + loader = WekaTraceLoader(run=run, prompt_generator=generator) + conversations = [] + for trace in eligible: + conversations.extend(loader.convert_to_conversations({trace.id: [trace]})) + if sum(len(c.turns) >= args.turns for c in conversations) >= args.streams: + break + inputs, sources = frozen_sessions(conversations, args.model, args.streams, args.turns) + from transformers import AutoTokenizer + + chat_tokenizer = AutoTokenizer.from_pretrained(str(args.model_directory), local_files_only=True) + for session, source in zip(inputs["data"], sources): + for payload, turn in zip(session["payloads"], source["turns"]): + count = prompt_length(chat_tokenizer, payload["messages"]) + if count + payload["max_tokens"] > 262144: + raise ValueError("Templated request exceeds context limit; never clip input or output") + turn["templated_prompt_tokens"] = count + args.output.mkdir(parents=True, exist_ok=False) + (args.output / "inputs.json").write_bytes(canonical(inputs) + b"\n") + (args.output / "sources.json").write_bytes(canonical(sources) + b"\n") + model_files = { + p.name: digest(p.read_bytes()) + for p in ( + args.model_directory / name + for name in ( + "config.json", + "model.safetensors.index.json", + "tokenizer.json", + "tokenizer_config.json", + "chat_template.jinja", + "generation_config.json", + "vocab.json", + "merges.txt", + "acquisition.json", + ) + ) + if p.is_file() + } + weights = sorted( + set(json.loads((args.model_directory / "model.safetensors.index.json").read_text())["weight_map"].values()) + ) + manifest = dict( + schema=1, + tool=identity, + dataset_revision=DATASET_REVISION, + source_sha256=digest(raw), + inputs_sha256=digest((args.output / "inputs.json").read_bytes()), + sources_sha256=digest((args.output / "sources.json").read_bytes()), + model_files=model_files, + weight_files={name: (args.model_directory / name).stat().st_size for name in weights}, + streams=args.streams, + turns_per_stream=args.turns, + requests=args.streams * args.turns, + output_tokens=sum(p["max_tokens"] for s in inputs["data"] for p in s["payloads"]), + scope="Fixed Weka-derived multi-turn control; frozen synthetic history; no think delays or DAG joins. " + "This is not the official AgentX scenario or a business-quality test.", + ) + (args.output / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + print( + json.dumps({k: manifest[k] for k in ("inputs_sha256", "sources_sha256", "requests", "output_tokens", "scope")}), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/prepare_aiperf_trace.py b/test/benchmark/qwen38/prepare_aiperf_trace.py new file mode 100644 index 0000000000..da8e294f2f --- /dev/null +++ b/test/benchmark/qwen38/prepare_aiperf_trace.py @@ -0,0 +1,72 @@ +import argparse +import hashlib +import json +from pathlib import Path +from urllib.request import urlopen + + +REVISION = "3cca71daccf2a7afb8fe3f0295358f70e3a69fdb" +URL = f"https://raw.githubusercontent.com/kvcache-ai/Mooncake/{REVISION}/FAST25-release/traces/conversation_trace.jsonl" +SHA256 = "b8cbb061a85206d729d91cdc2981f43c9e0d99209dce588d3af5f7934408b9df" + + +def select_requests(rows, count, context_length): + if count < 1 or context_length < 2: + raise ValueError("Positive count and context length >= 2 are required") + eligible = [] + previous = -1 + for index, row in enumerate(rows): + if row["timestamp"] < previous: + raise ValueError("Trace timestamps must be nondecreasing") + previous = row["timestamp"] + if ( + row["input_length"] > 0 + and row["output_length"] > 0 + and row["input_length"] + row["output_length"] <= context_length + ): + eligible.append((index, row)) + if len(eligible) < count: + raise ValueError("Not enough requests fit the shared context limit") + return eligible[:count], len(eligible) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-directory", type=Path, required=True) + parser.add_argument("--source", type=Path) + parser.add_argument("--count", type=int, default=256) + parser.add_argument("--context-length", type=int, default=4096) + args = parser.parse_args() + raw = args.source.read_bytes() if args.source else urlopen(URL, timeout=60).read() + if hashlib.sha256(raw).hexdigest() != SHA256: + raise ValueError("Public trace differs from the pinned source") + rows = [json.loads(line) for line in raw.splitlines() if line.strip()] + selected, eligible = select_requests(rows, args.count, args.context_length) + directory = args.output_directory + directory.mkdir(parents=True, exist_ok=True) + (directory / "conversation_trace.original.jsonl").write_bytes(raw) + trace = directory / "replay.jsonl" + trace.write_text("".join(json.dumps(row, separators=(",", ":")) + "\n" for _, row in selected)) + report = dict( + source_url=URL, + revision=REVISION, + source_sha256=SHA256, + source_requests=len(rows), + eligible_requests=eligible, + selected_indices=[index for index, _ in selected], + requests=len(selected), + context_length=args.context_length, + selection="First eligible requests in source order; no input/output clipping", + prompt_content="Anonymized lengths and 512-token prefix hashes; AIPerf synthesizes text", + timing="Original millisecond timestamps, normalized by AIPerf; speed multipliers reported separately", + trace_sha256=hashlib.sha256(trace.read_bytes()).hexdigest(), + schedule_span_ms=selected[-1][1]["timestamp"] - selected[0][1]["timestamp"], + input_tokens=sum(row["input_length"] for _, row in selected), + output_tokens=sum(row["output_length"] for _, row in selected), + ) + (directory / "manifest.json").write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/retrieval_quality.py b/test/benchmark/qwen38/retrieval_quality.py new file mode 100644 index 0000000000..1206964ff1 --- /dev/null +++ b/test/benchmark/qwen38/retrieval_quality.py @@ -0,0 +1,51 @@ +import hashlib +import json +import re + +from test.benchmark.qwen38.check_generation import REVISION + +MANIFEST_SHA256 = "3214c9f09a0eb3bae49061d4658daeaea251b4befcfa5214278d796e5762d7cd" +LABELS = ("ALPHA", "BETA", "GAMMA") +ANSWER = re.compile(r"\s*ALPHA\s*=\s*(\d+)\s*;\s*BETA\s*=\s*(\d+)\s*;\s*GAMMA\s*=\s*(\d+)\s*") + + +def score_records(text, expected): + if set(expected) != set(LABELS) or any( + not isinstance(value, str) or not re.fullmatch(r"[0-9]+", value) for value in expected.values() + ): + raise ValueError("Require the three declared record labels and exact digit strings") + match = ANSWER.fullmatch(text) if isinstance(text, str) else None + predicted = dict(zip(LABELS, match.groups())) if match else None + records = {label: predicted is not None and predicted[label] == expected[label] for label in LABELS} + return dict(predicted=predicted, expected=expected, records=records, correct=all(records.values())) + + +def load_workload(path, model_directory): + manifest = json.loads(path.read_text()) + if hashlib.sha256(json.dumps(manifest, sort_keys=True).encode()).hexdigest() != MANIFEST_SHA256: + raise ValueError("Retrieval workload differs from pinned preparation") + config = json.loads((model_directory / "config.json").read_text()) + context = config.get("text_config", config)["max_position_embeddings"] + if ( + manifest["status"] != "prepared" + or manifest["model_revision"] != REVISION + or manifest["context_limit"] != context + or manifest["decoding"] != dict(temperature=0, ignore_eos=False, enable_thinking=False) + ): + raise ValueError("Retrieval workload model, context or decoding differs") + for name in ("tokenizer.json", "tokenizer_config.json", "chat_template.jinja"): + if hashlib.sha256((model_directory / name).read_bytes()).hexdigest() != manifest["tokenizer_files"].get(name): + raise ValueError("Tokenizer differs from retrieval preparation") + for case in manifest["cases"]: + ids, limit = case["prompt_ids"], case["max_output_tokens"] + if ( + not isinstance(ids, list) + or not ids + or any(type(token) is not int or token < 0 for token in ids) + or type(limit) is not int + or limit < 1 + or len(ids) + limit > context + ): + raise ValueError("Retrieval inputs exceed context or lack exact token IDs") + score_records("", case["expected_records"]) + return manifest diff --git a/test/benchmark/qwen38/run_agentx.py b/test/benchmark/qwen38/run_agentx.py new file mode 100644 index 0000000000..297b34adea --- /dev/null +++ b/test/benchmark/qwen38/run_agentx.py @@ -0,0 +1,267 @@ +import argparse +import datetime +import json +import os +import signal +import time +from pathlib import Path +import subprocess +import sys +import urllib.request + +from test.benchmark.qwen38.audit_agentx import audit, repeatability +from test.benchmark.qwen38.prepare_agentx import canonical, digest, tool_identity + + +TASK = "a3844d54-8a2b-4532-bf54-f657c9ca72a0" + + +def http(base, path, method="GET"): + opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) + with opener.open(urllib.request.Request(base + path, method=method), timeout=120) as response: + return response.read() + + +def container_info(name): + result = subprocess.check_output(["sudo", "-n", "docker", "inspect", name], text=True) + info = json.loads(result)[0] + if not info["State"]["Running"] or info["Config"].get("Labels", {}).get("qwen38.task") != TASK: + raise ValueError("Require a running container owned by this task") + others = subprocess.check_output(["sudo", "-n", "docker", "ps", "-q"], text=True).split() + if others != [info["Id"][:12]]: + raise ValueError("Unexpected concurrent container; run on the isolated benchmark host") + return info + + +def idle_window(container, history, timeout=300): + samples = [] + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + owners = subprocess.check_output( + ["sudo", "-n", "docker", "top", container, "-eo", "pid"], text=True + ).splitlines()[1:] + allowed = {p.strip() for p in owners} + processes = subprocess.check_output( + ["nvidia-smi", "--query-compute-apps=pid", "--format=csv,noheader,nounits"], text=True + ).split() + if set(processes) - allowed: + raise ValueError("GPU process outside the benchmark container") + raw = subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=index,utilization.gpu,temperature.gpu", + "--format=csv,noheader,nounits", + ], + text=True, + ) + rows = [[int(v.strip()) for v in line.split(",")] for line in raw.splitlines()] + if sorted(index for index, _, _ in rows) != list(range(8)): + raise ValueError("Require all eight benchmark GPUs") + history.append(dict(timestamp=datetime.datetime.now().isoformat(), gpus=rows)) + if any(util > 5 or temperature > 60 for _, util, temperature in rows): + samples.clear() + else: + samples.append(rows) + time.sleep(10) + if len(samples) == 3: + return samples + raise TimeoutError(f"GPUs did not stay idle (utilization <=5%, temperature <=60C) within {timeout} seconds") + + +def reset_cache(base, engine): + path = "/reset_prefix_cache" if engine == "vllm" else "/flush_cache" + raw = http(base, path, "POST") + if engine in ("vllm", "lightllm"): + if json.loads(raw).get("success") is not True: + raise ValueError("Server did not confirm cache reset") + else: + confirmation = raw.strip().lower() + if confirmation not in (b"true", b"", b"null") and not confirmation.startswith(b"cache flushed.\n"): + raise ValueError("Unexpected SGLang cache reset response") + return dict(path=path, response=raw.decode()) + + +def fixed_config(base, model, tokenizer, inputs, count, streams, output): + return { + "schemaVersion": "2.0", + "randomSeed": 42, + "benchmark": dict( + model=model, + endpoint=dict( + url=base, + type="chat", + streaming=True, + use_server_token_count=True, + timeout=1200, + wait_for_model_timeout=0, + connection_reuse="pooled", + ), + tokenizer=dict(name=str(tokenizer)), + dataset=dict(type="file", format="inputs_json", path=str(inputs), sampling="sequential"), + phases=[dict(name="profiling", type="concurrency", concurrency=streams, sessions=count)], + artifacts=dict( + dir=str(output), + summary=["json"], + records=["jsonl"], + raw=True, + export_outputs_json=False, + auto_plot=False, + ), + runtime=dict(ui="none", stats_interval=30), + gpu_telemetry=dict(enabled=False), + ), + } + + +def execute(command, log, environment): + with log.open("w") as handle: + process = subprocess.Popen(command, stdout=handle, stderr=subprocess.STDOUT, env=environment) + log.with_suffix(".pid").write_text(str(process.pid) + "\n") + try: + returncode = process.wait() + except BaseException: + process.terminate() + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=10) + raise + if returncode: + raise RuntimeError(f"AgentX exited {returncode}; inspect {log}") + + +def main(): + parser = argparse.ArgumentParser(description="Run three independently warmed fixed AgentX trials") + parser.add_argument("--engine", choices=["vllm", "sglang", "lightllm"], required=True) + parser.add_argument("--container", required=True) + parser.add_argument("--workload", type=Path, required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--model", default="Qwen3.8-Flash-Next") + parser.add_argument("--url", default="http://127.0.0.1:28081") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--repeats", type=int, default=3) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + if args.repeats < 3 or args.repeats > 10: + parser.error("Require three to ten trials") + if args.url != "http://127.0.0.1:28081": + parser.error("This protocol is restricted to the isolated local benchmark endpoint") + args.output = args.output.resolve() + args.workload = args.workload.resolve() + manifest = json.loads((args.workload / "manifest.json").read_text()) + if canonical(tool_identity()) != canonical(manifest["tool"]): + raise ValueError("Installed tool/dependencies changed from the frozen manifest") + inputs_path = args.workload / "inputs.json" + if digest(inputs_path.read_bytes()) != manifest["inputs_sha256"]: + raise ValueError("Frozen workload changed") + if digest((args.workload / "sources.json").read_bytes()) != manifest["sources_sha256"]: + raise ValueError("Frozen workload provenance changed") + for name, expected in manifest["model_files"].items(): + if digest((args.model_directory / name).read_bytes()) != expected: + raise ValueError(f"Checkpoint/tokenizer changed: {name}") + for name, size in manifest["weight_files"].items(): + if (args.model_directory / name).stat().st_size != size: + raise ValueError(f"Checkpoint shard size changed: {name}") + info = container_info(args.container) + args.output.mkdir(parents=True, exist_ok=False) + inputs = json.loads(inputs_path.read_text()) + report = dict( + status="running", + mode="fixed", + engine=args.engine, + manifest=manifest, + container=info, + runs=[], + scope="Fixed Weka-derived control; not the official response-dependent AgentX scenario", + protocol_files={p.name: digest(p.read_bytes()) for p in sorted(Path(__file__).parent.glob("*_agentx.py"))}, + cache_policy="Confirmed reset before every complete warmup; same full fixed-work warmup every trial; " + "reset again before measurement.", + ) + (args.output / "hardware.txt").write_text(subprocess.check_output(["nvidia-smi", "-q"], text=True)) + environment = os.environ.copy() + environment.update( + NO_PROXY="127.0.0.1,localhost", + no_proxy="127.0.0.1,localhost", + HF_HUB_OFFLINE="1", + TRANSFORMERS_OFFLINE="1", + PYTHONHASHSEED="42", + AIPERF_HTTP_KEEPALIVE_TIMEOUT="4", + ) + entry = "from aiperf.cli import app; app()" + + def save(): + (args.output / "suite.json").write_text(json.dumps(report, indent=2) + "\n") + + monitor = None + + def interrupted(signum, frame): + raise KeyboardInterrupt(f"Signal {signum}") + + signal.signal(signal.SIGTERM, interrupted) + try: + save() + gpu_log = (args.output / "gpu-samples.csv").open("w") + monitor = subprocess.Popen( + [ + "nvidia-smi", + "--query-gpu=timestamp,index,utilization.gpu,memory.used,power.draw,temperature.gpu", + "--format=csv", + "--loop=5", + ], + stdout=gpu_log, + ) + for index in range(args.repeats): + if container_info(args.container)["Id"] != info["Id"]: + raise ValueError("Server container changed between trials") + trial = args.output / f"trial-{index + 1:02}" + trial.mkdir() + result = dict(trial=index + 1) + report["runs"].append(result) + save() + result["reset_before_warmup"] = reset_cache(args.url, args.engine) + for phase in ("warmup", "measure"): + if phase == "measure": + result["reset_before_measure"] = reset_cache(args.url, args.engine) + result[phase + "_idle_checks"] = [] + result[phase + "_idle_samples"] = idle_window(args.container, result[phase + "_idle_checks"]) + save() + output = trial / phase + config = fixed_config( + args.url, args.model, args.model_directory, inputs_path, manifest["streams"], 32, output + ) + config_path = trial / (phase + ".json") + config_path.write_text(json.dumps(config, indent=2) + "\n") + command = [sys.executable, "-c", entry, "profile", "--config", str(config_path)] + (trial / (phase + "-command.json")).write_text(json.dumps(command, indent=2) + "\n") + print(datetime.datetime.now().isoformat(), index + 1, phase, output, flush=True) + (trial / (phase + "-metrics-before.txt")).write_bytes(http(args.url, "/metrics")) + execute(command, trial / (phase + ".log"), environment) + if container_info(args.container)["Id"] != info["Id"]: + raise ValueError("Container changed during trial") + (trial / (phase + "-metrics-after.txt")).write_bytes(http(args.url, "/metrics")) + result[phase + "_audit"] = audit(output, inputs) + save() + result["audit"] = result["measure_audit"] + save() + report["repeatability"] = repeatability([r["audit"] for r in report["runs"]]) + report["status"] = report["repeatability"]["status"] + except BaseException as error: + report.update(status="failed", error=repr(error)) + raise + finally: + if monitor is not None: + monitor.terminate() + monitor.wait(timeout=20) + gpu_log.close() + logs = subprocess.run(["sudo", "-n", "docker", "logs", args.container], capture_output=True, text=True) + (args.output / "server.log").write_text(logs.stdout + logs.stderr) + save() + print(json.dumps(report.get("repeatability", {"status": report["status"]}), indent=2), flush=True) + if report["status"] == "nonrepeatable": + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/score_reference.py b/test/benchmark/qwen38/score_reference.py new file mode 100644 index 0000000000..a42d9b66c2 --- /dev/null +++ b/test/benchmark/qwen38/score_reference.py @@ -0,0 +1,208 @@ +import argparse +import hashlib +import importlib.metadata +import importlib.util +import json +import os +from pathlib import Path +import subprocess +import sys +import time + +import torch +from safetensors import safe_open + +from check_generation import REVISION, cases_for, verified_checkpoint + +SOURCE_SHA256 = "87c27c88424cea73d5d8d9179f76a77a97e6e4249ce43ecff79a8b255fab0504" + + +def load_source(path): + if hashlib.sha256(path.read_bytes()).hexdigest() != SOURCE_SHA256: + raise ValueError("Reference equations do not match the pinned source") + name = "transformers.models.qwen4_exp.modeling_qwen4_exp" + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def verify_embedding_rows(model, root): + index = json.loads((root / "model.safetensors.index.json").read_text())["weight_map"] + weight = model.model.language_model.layers[1].ple.ple_embedding.ngram_embedding.weight + if weight.is_meta or tuple(weight.shape) != (320001536, 160) or weight.dtype != torch.bfloat16: + raise ValueError("Reference lexical table has the wrong shape, residency or dtype") + checked = 0 + for shard in range(128): + key = f"model.language_model.layers.1.ple.ple_embedding.ngram_embedding.shard_{shard}.weight" + with safe_open(root / index[key], framework="pt", device="cpu") as source: + tensor = source.get_slice(key) + rows, width = tensor.get_shape() + if (rows, width) != (2500012, 160): + raise ValueError(f"Unexpected source shard shape: {key}") + for row in (0, rows // 2, rows - 1): + actual = weight[shard * rows + row].cpu() + torch.testing.assert_close(actual, tensor[row : row + 1][0], atol=0, rtol=0) + checked += 1 + return {"shards": 128, "exact_rows": checked, "device": str(weight.device)} + + +def forward_reference(model, inputs, cache, length, chunk_size=None): + count = inputs.shape[1] + if count < 1 or length < count or (chunk_size is not None and chunk_size < 1): + raise ValueError("Reference forward requires nonempty inputs and valid lengths") + chunk_size = count if chunk_size is None else chunk_size + for start in range(0, count, chunk_size): + stop = min(start + chunk_size, count) + absolute_start, absolute_stop = length - count + start, length - count + stop + output = model( + input_ids=inputs[:, start:stop], + position_ids=torch.arange(absolute_start, absolute_stop, device=inputs.device).unsqueeze(0), + attention_mask=torch.ones((1, absolute_stop), device=inputs.device, dtype=torch.long), + past_key_values=cache, + use_cache=True, + logits_to_keep=1, + ) + cache = output.past_key_values + return output + + +@torch.inference_mode() +def score_case(model, prompt, result, directory, case, prefill_chunk_size=None): + tokens = result["token_ids"] + native = result["response"]["tokens"] + if [item["id"] for item in native] != tokens: + raise ValueError("Native probability records do not match generated IDs") + device = model.model.language_model.embed_tokens.weight.device + inputs = torch.tensor([prompt], device=device, dtype=torch.long) + cache = None + rows, scores = [], [] + for step, target in enumerate(tokens): + length = len(prompt) + step + output = forward_reference(model, inputs, cache, length, prefill_chunk_size) + cache = output.past_key_values + logits = output.logits[0, -1].float().cpu() + if not bool(torch.isfinite(logits).all()): + raise ValueError(f"Nonfinite reference logits at {case} step {step}") + scores.append(logits) + reference_logprob = float(logits[target] - logits.logsumexp(0)) + native_logprob = float(native[step]["logprob"]) + greedy = int(logits.argmax()) + rows.append( + dict( + step=step, + token_id=target, + reference_greedy_id=greedy, + reference_logprob=reference_logprob, + native_logprob=native_logprob, + logprob_delta=native_logprob - reference_logprob, + reference_greedy_gap=float(logits[greedy] - logits[target]), + ) + ) + inputs = torch.tensor([[target]], device=device, dtype=torch.long) + torch.save(torch.stack(scores), directory / f"{case}-scores.pt") + return dict( + steps=rows, + greedy_matches=sum(row["token_id"] == row["reference_greedy_id"] for row in rows), + reference_mean_nll=-sum(row["reference_logprob"] for row in rows) / len(rows), + native_mean_nll=-sum(row["native_logprob"] for row in rows) / len(rows), + mean_absolute_logprob_delta=sum(abs(row["logprob_delta"]) for row in rows) / len(rows), + maximum_absolute_logprob_delta=max(abs(row["logprob_delta"]) for row in rows), + ) + + +def main(): + if not os.environ.get("EXP_ACTIVE") or os.environ.get("USE_HUB_KERNELS") != "NO": + raise RuntimeError("Run through exp with USE_HUB_KERNELS=NO") + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--source", type=Path, required=True) + parser.add_argument("--device-plan", type=Path, required=True) + parser.add_argument("--generation", type=Path, required=True) + parser.add_argument("--cases", nargs="+", default=["english", "chinese", "code_unicode"]) + parser.add_argument("--prefill-chunk-size", type=int) + args = parser.parse_args() + if args.prefill_chunk_size is not None and args.prefill_chunk_size < 1: + parser.error("Prefill chunk size must be positive") + root, run = args.model_directory, Path(os.environ["EXP_RUN_DIR"]) + verified_checkpoint(root) + trace = json.loads(args.generation.read_text()) + cases = cases_for(root) + if trace["revision"] != REVISION or trace["cases"] != cases or trace["service"]["framework"] != "lightllm": + raise ValueError("Generation must contain native outputs for the pinned shared inputs") + for case in args.cases: + if not trace["results"][case]["token_ids"]: + raise ValueError(f"Missing generated IDs: {case}") + plan = json.loads(args.device_plan.read_text()) + device_map = plan["device_map"] + if not device_map or not all(type(device) is int for device in device_map.values()): + raise ValueError("This reference diagnostic requires a complete GPU placement") + module = load_source(args.source) + report = dict( + status="loading", + scope="identical native token histories; diagnostic only, no quality parity or performance claim", + revision=REVISION, + source_sha256=SOURCE_SHA256, + source=str(args.source.resolve()), + generation=str(args.generation.resolve()), + device_plan=plan, + command=sys.argv, + cuda_visible_devices=os.environ.get("CUDA_VISIBLE_DEVICES"), + software={name: importlib.metadata.version(name) for name in ("torch", "transformers", "accelerate")}, + hardware=subprocess.check_output( + ["nvidia-smi", "--query-gpu=index,name,uuid,memory.used,memory.free", "--format=csv,noheader"], text=True + ), + configuration=dict( + dtype="bfloat16", + attention="sdpa", + experts="eager", + use_hub_kernels=False, + allow_tf32=torch.backends.cuda.matmul.allow_tf32, + bf16_reduced_precision=torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction, + deterministic_algorithms=torch.are_deterministic_algorithms_enabled(), + prefill_chunk_size=args.prefill_chunk_size, + ), + results={}, + ) + path = run / "reference-scores.json" + path.write_text(json.dumps(report, indent=2)) + started = time.monotonic() + try: + model, loading = module.Qwen4ExpForConditionalGeneration.from_pretrained( + str(root), + dtype=torch.bfloat16, + device_map=device_map, + local_files_only=True, + use_safetensors=True, + attn_implementation="sdpa", + experts_implementation="eager", + output_loading_info=True, + ) + model.eval() + (run / "loading-info.json").write_text(json.dumps(loading, indent=2, default=str)) + if loading.get("missing_keys") or loading.get("mismatched_keys"): + raise ValueError("Reference checkpoint loading is incomplete") + if any(not key.startswith("mtp.") for key in loading.get("unexpected_keys", [])): + raise ValueError("Unexpected non-MTP reference weights") + report["lexical_row_verification"] = verify_embedding_rows(model, root) + report.update(status="scoring", load_seconds=time.monotonic() - started) + path.write_text(json.dumps(report, indent=2)) + for case in args.cases: + result = score_case(model, cases[case], trace["results"][case], run, case, args.prefill_chunk_size) + report["results"][case] = result + path.write_text(json.dumps(report, indent=2)) + print( + json.dumps({"case": case, **{key: value for key, value in result.items() if key != "steps"}}), + flush=True, + ) + report["status"] = "recorded" + except Exception as error: + report.update(status="failed", error_type=type(error).__name__, error=str(error)) + raise + finally: + path.write_text(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/service/cache_capacity.py b/test/benchmark/service/cache_capacity.py new file mode 100644 index 0000000000..0fb53a7578 --- /dev/null +++ b/test/benchmark/service/cache_capacity.py @@ -0,0 +1,50 @@ +import math + +from prometheus_client.parser import text_string_to_metric_families + + +def plan_vllm_capacity(metrics_text, model_length, target_tokens): + if any(type(value) is not int or value < 1 for value in (model_length, target_tokens)): + raise ValueError("Require positive integer context and token budgets") + if target_tokens % model_length: + raise ValueError("Token budget must contain whole full-context requests") + samples = [ + sample + for family in text_string_to_metric_families(metrics_text) + for sample in family.samples + if sample.name == "vllm:cache_config_info" + ] + if len(samples) != 1 or samples[0].value != 1: + raise ValueError("Require exactly one initialized cache-info sample") + fields = samples[0].labels + try: + blocks = int(fields["num_gpu_blocks"]) + concurrency = float(fields["kv_cache_max_concurrency"]) + reported_tokens = int(fields["kv_cache_size_tokens"]) + except (KeyError, TypeError, ValueError) as error: + raise ValueError("Initialized block count and group-aware capacity are required") from error + if blocks < 1 or not math.isfinite(concurrency) or concurrency <= 0 or reported_tokens < 1: + raise ValueError("Runtime cache geometry must be finite and positive") + if reported_tokens != int(concurrency * model_length): + raise ValueError("Reported capacity does not match the recorded context length") + if target_tokens > reported_tokens: + raise ValueError("Capacity planning only shrinks the observed initialized pool") + ratio = blocks / concurrency + per_request = round(ratio) + if per_request < 1 or not math.isclose(ratio, per_request, rel_tol=1e-12, abs_tol=1e-9): + raise ValueError("Cannot infer an integer device-pool block budget per request") + return dict( + status="planned", + scope="Logical single-engine budget inferred from initialized metrics; verify again after launch", + model_length=model_length, + target_tokens=target_tokens, + observed_tokens=reported_tokens, + observed_num_gpu_blocks=blocks, + observed_max_concurrency=concurrency, + blocks_per_full_request=per_request, + num_gpu_blocks_override=per_request * (target_tokens // model_length), + cache_geometry={ + key: fields.get(key) + for key in ("block_size", "mamba_block_size", "mamba_cache_mode", "mamba_ssm_cache_dtype") + }, + ) diff --git a/test/benchmark/service/completion.py b/test/benchmark/service/completion.py new file mode 100644 index 0000000000..9cb190a53c --- /dev/null +++ b/test/benchmark/service/completion.py @@ -0,0 +1,61 @@ +def completion_request(framework, model, prompt_ids, limit): + if not prompt_ids or any(type(token) is not int or token < 0 for token in prompt_ids): + raise ValueError("Require nonnegative integer prompt IDs") + if type(limit) is not int or limit < 1: + raise ValueError("Require a positive integer output limit") + if framework == "lightllm": + return "/generate", { + "inputs": prompt_ids, + "parameters": dict( + do_sample=False, + ignore_eos=False, + max_new_tokens=limit, + return_details=True, + print_eos_token=True, + ), + } + if framework in ("vllm", "sglang"): + return "/v1/completions", dict( + model=model, + prompt=prompt_ids, + max_tokens=limit, + temperature=0, + seed=3107, + ignore_eos=False, + return_token_ids=True, + ) + raise ValueError("Unsupported framework") + + +def completion_tokens(body, framework, prompt_ids, limit, eos_ids): + if framework == "lightllm": + ids = [token["id"] for token in body["tokens"]] + prompt_count, output_count = body["prompt_tokens"], body["count_output_tokens"] + reason = body["finish_reason"] + elif framework in ("vllm", "sglang"): + if len(body["choices"]) != 1 or body["choices"][0].get("index") != 0: + raise ValueError("Require exactly one completion choice") + choice = body["choices"][0] + if choice.get("prompt_token_ids", prompt_ids) != prompt_ids: + raise ValueError("Service changed the exact prompt IDs") + ids, reason = choice["token_ids"], choice["finish_reason"] + prompt_count, output_count = body["usage"]["prompt_tokens"], body["usage"]["completion_tokens"] + else: + raise ValueError("Unsupported framework") + if not isinstance(ids, list) or not ids or any(type(token) is not int or token < 0 for token in ids): + raise ValueError("Require actual nonnegative output token IDs") + if ( + type(prompt_count) is not int + or type(output_count) is not int + or prompt_count != len(prompt_ids) + or output_count != len(ids) + or len(ids) > limit + ): + raise ValueError("Service counts differ from exact prompt/output IDs") + if reason not in ("stop", "length") or any(token in eos_ids for token in ids[:-1]): + raise ValueError("Invalid termination reason or output beyond first EOS") + if reason == "length" and (len(ids) != limit or ids[-1] in eos_ids): + raise ValueError("Length termination must exhaust the budget without EOS") + if framework == "lightllm" and reason == "stop" and ids[-1] not in eos_ids: + raise ValueError("Native natural termination must expose its declared EOS") + return dict(token_ids=ids, finish_reason=reason, truncated=reason == "length") diff --git a/test/benchmark/service/completion_runner.py b/test/benchmark/service/completion_runner.py new file mode 100644 index 0000000000..95a7a56cdd --- /dev/null +++ b/test/benchmark/service/completion_runner.py @@ -0,0 +1,63 @@ +from concurrent.futures import ThreadPoolExecutor +import threading + +import requests + +from .completion import completion_request, completion_tokens + + +class CompletionClient: + def __init__( + self, + framework, + model, + base_url, + timeout, + eos_ids, + raw_directory, + *, + concurrency=1, + session_factory=requests.Session, + ): + if type(concurrency) is not int or not 1 <= concurrency <= 8: + raise ValueError("Completion concurrency must be between one and eight") + self.framework, self.model, self.base_url = framework, model, base_url + self.timeout, self.eos_ids, self.raw_directory = timeout, eos_ids, raw_directory + self.concurrency, self.session_factory = concurrency, session_factory + self.local = threading.local() + self.sessions, self.lock, self.executor = [], threading.Lock(), None + + def _complete(self, case): + index = case["dataset_index"] + if type(index) is not int or index < 0: + raise ValueError("Require a nonnegative integer dataset index") + if not hasattr(self.local, "session"): + session = self.session_factory() + session.trust_env = False + self.local.session = session + with self.lock: + self.sessions.append(session) + endpoint, payload = completion_request( + self.framework, self.model, case["prompt_ids"], case["max_output_tokens"] + ) + response = self.local.session.post(self.base_url.rstrip("/") + endpoint, json=payload, timeout=self.timeout) + (self.raw_directory / f"response-{index}.txt").write_text(response.text) + response.raise_for_status() + return completion_tokens( + response.json(), self.framework, case["prompt_ids"], case["max_output_tokens"], self.eos_ids + ) + + def results(self, cases): + if self.concurrency == 1: + return map(self._complete, cases) + if self.executor is not None: + raise RuntimeError("A completion client accepts one concurrent workload") + self.executor = ThreadPoolExecutor(max_workers=self.concurrency) + + return self.executor.map(self._complete, cases) + + def close(self): + if self.executor is not None: + self.executor.shutdown(wait=True, cancel_futures=True) + for session in self.sessions: + session.close() diff --git a/test/benchmark/service/stream_cohort.py b/test/benchmark/service/stream_cohort.py new file mode 100644 index 0000000000..9ebc722ddc --- /dev/null +++ b/test/benchmark/service/stream_cohort.py @@ -0,0 +1,166 @@ +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +import math +from pathlib import Path +from queue import Empty, SimpleQueue +from statistics import fmean +import threading + + +@dataclass(frozen=True) +class StreamingWorkload: + name: str + prompt_ids: tuple[int, ...] + expected_ids: tuple[int, ...] + + def __post_init__(self): + if not isinstance(self.name, str) or not self.name: + raise ValueError("Workload requires a name") + for field in ("prompt_ids", "expected_ids"): + values = getattr(self, field) + if not isinstance(values, (list, tuple)) or not values or any(type(x) is not int or x < 0 for x in values): + raise ValueError("Workload requires exact nonnegative prompt and reference token IDs") + object.__setattr__(self, field, tuple(values)) + + +def distribution(values): + if not values: + return None + if any(not math.isfinite(value) or value < 0 for value in values): + raise ValueError("Latency samples must be finite and nonnegative") + ordered = sorted(values) + result = dict(count=len(values), min=ordered[0], max=ordered[-1], mean=fmean(values)) + for percentile in (50, 90, 95, 99): + position = (len(ordered) - 1) * percentile / 100 + left, right = math.floor(position), math.ceil(position) + result[f"p{percentile}"] = ordered[left] + (ordered[right] - ordered[left]) * (position - left) + return result + + +def _validate_result(result): + metrics, timeline = result["metrics"], result["timeline"] + started, first, last, ended = (timeline[key] for key in ("started", "first_token", "last_token", "ended")) + if not all(math.isfinite(value) for value in (started, first, last, ended)): + raise ValueError("Request timeline must be finite") + if not started <= first <= last <= ended or ended <= started: + raise ValueError("Request timeline is incomplete or nonmonotonic") + ids = result["token_ids"] + if not ids or any(type(token) is not int or token < 0 for token in ids): + raise ValueError("Result requires actual integer output IDs") + if type(metrics["input_tokens"]) is not int or metrics["input_tokens"] < 1: + raise ValueError("Result requires the actual input token count") + if type(metrics["output_tokens"]) is not int or metrics["output_tokens"] != len(ids): + raise ValueError("Output count differs from actual token IDs") + gaps = metrics["receiver_itl_s"] + if len(gaps) != len(ids) - 1 or any(not math.isfinite(gap) or gap < 0 for gap in gaps): + raise ValueError("Receiver gaps must account for every output position") + expected = dict(ttft_s=first - started, e2e_s=ended - started) + if len(ids) > 1: + expected["tpot_s"] = (last - first) / (len(ids) - 1) + elif metrics["tpot_s"] is not None: + raise ValueError("A single output token has no TPOT") + for key, value in expected.items(): + if not math.isclose(metrics[key], value, rel_tol=1e-9, abs_tol=1e-9): + raise ValueError(f"{key} differs from the captured request timeline") + if not math.isclose(sum(gaps), last - first, rel_tol=1e-9, abs_tol=1e-9): + raise ValueError("Receiver gaps differ from the token arrival interval") + + +def summarize_cohort(results): + if not results: + raise ValueError("An empty cohort cannot be measured") + for result in results: + _validate_result(result) + elapsed = max(item["timeline"]["ended"] for item in results) - min(item["timeline"]["started"] for item in results) + input_tokens = sum(item["metrics"]["input_tokens"] for item in results) + output_tokens = sum(len(item["token_ids"]) for item in results) + cache_counts = [item["metrics"].get("cached_input_tokens") for item in results] + if any( + count is not None and (type(count) is not int or not 0 <= count <= item["metrics"]["input_tokens"]) + for count, item in zip(cache_counts, results) + ): + raise ValueError("Invalid cached input count") + cached_tokens = sum(cache_counts) if all(count is not None for count in cache_counts) else None + events = sorted( + [(item["timeline"]["started"], 1) for item in results] + [(item["timeline"]["ended"], -1) for item in results] + ) + active = maximum = 0 + for _, delta in events: + active += delta + maximum = max(maximum, active) + gaps = [gap for item in results for gap in item["metrics"]["receiver_itl_s"]] + return dict( + requests=len(results), + input_tokens=input_tokens, + output_tokens=output_tokens, + cache_reported_requests=sum(count is not None for count in cache_counts), + cached_input_tokens=cached_tokens, + uncached_input_tokens=input_tokens - cached_tokens if cached_tokens is not None else None, + elapsed_s=elapsed, + requests_per_s=len(results) / elapsed, + input_tokens_per_s=input_tokens / elapsed, + output_tokens_per_s=output_tokens / elapsed, + input_token_definition="requested prompt IDs; includes cached tokens", + maximum_client_inflight=maximum, + ttft_s=distribution([item["metrics"]["ttft_s"] for item in results]), + tpot_s=distribution([item["metrics"]["tpot_s"] for item in results if item["metrics"]["tpot_s"] is not None]), + e2e_s=distribution([item["metrics"]["e2e_s"] for item in results]), + receiver_itl_s=distribution(gaps), + receiver_zero_gap_fraction=sum(gap == 0 for gap in gaps) / len(gaps) if gaps else None, + itl_definition="client arrival gaps; tokens delivered together share a timestamp", + percentile_definition="linear interpolation between ordered observations", + ) + + +def run_cohort(workloads, concurrency, request, response_directory: Path): + workloads = tuple(workloads) + if type(concurrency) is not int or not 1 <= concurrency <= len(workloads): + raise ValueError("Concurrency must be positive and no larger than the request count") + if any(not isinstance(item, StreamingWorkload) for item in workloads): + raise ValueError("Every request requires a qualified streaming workload") + response_directory.mkdir(parents=True, exist_ok=True) + pending = SimpleQueue() + for index, workload in enumerate(workloads): + pending.put((index, workload)) + records = [None] * len(workloads) + ready = threading.Barrier(concurrency + 1) + release = threading.Event() + + def worker(): + ready.wait(timeout=30) + release.wait() + while True: + try: + index, workload = pending.get_nowait() + except Empty: + return + path = response_directory / f"{index:06d}.sse" + record = dict(index=index, workload=workload.name, response=str(path), status="running") + try: + result = request(workload, path) + record["result"] = result + _validate_result(result) + if result["metrics"]["input_tokens"] != len(workload.prompt_ids): + raise ValueError("Service input length differs from the qualified workload") + if tuple(result["token_ids"]) != workload.expected_ids: + raise ValueError("Output IDs differ from the qualified reference") + record["status"] = "passed" + except Exception as error: + record.update(status="failed", error_type=type(error).__name__, error=str(error)) + records[index] = record + + with ThreadPoolExecutor(max_workers=concurrency) as executor: + futures = [executor.submit(worker) for _ in range(concurrency)] + try: + ready.wait(timeout=30) + finally: + release.set() + for future in futures: + future.result() + passed = all(item["status"] == "passed" for item in records) + return dict( + status="passed" if passed else "failed", + concurrency=concurrency, + records=records, + metrics=summarize_cohort([item["result"] for item in records]) if passed else None, + ) diff --git a/test/benchmark/service/stream_timing.py b/test/benchmark/service/stream_timing.py new file mode 100644 index 0000000000..3bce3d7254 --- /dev/null +++ b/test/benchmark/service/stream_timing.py @@ -0,0 +1,80 @@ +import json +import math + + +def _stats(values): + if not values: + return None + ordered = sorted(values) + + def percentile(fraction): + index = (len(ordered) - 1) * fraction + lower, upper = math.floor(index), math.ceil(index) + return ordered[lower] + (ordered[upper] - ordered[lower]) * (index - lower) + + return dict(count=len(ordered), p50=percentile(0.5), p95=percentile(0.95), p99=percentile(0.99), max=ordered[-1]) + + +def _nanoseconds(value): + if type(value) is not int or value < 0: + raise ValueError("Expected a nonnegative integer timestamp in nanoseconds") + return value + + +def streaming_completion_times(record): + metadata = record["metadata"] + if record["status"] != 200 or metadata.get("was_cancelled"): + raise ValueError("Cannot measure a failed or cancelled streaming response") + start = _nanoseconds(metadata["request_start_ns"]) + end = _nanoseconds(metadata["request_end_ns"]) + origin = _nanoseconds(record["start_perf_ns"]) + if end < start: + raise ValueError("HTTP end precedes request start") + text_times, done_times = [], [] + previous = origin + for response in record["responses"]: + timestamp = _nanoseconds(response["perf_ns"]) + if timestamp < previous: + raise ValueError("Nonmonotonic streaming packet timestamps") + previous = timestamp + for packet in response.get("packets", []): + if packet.get("name") != "data": + continue + if done_times: + raise ValueError("Unexpected SSE data after DONE") + if packet["value"] == "[DONE]": + done_times.append(timestamp) + continue + payload = json.loads(packet["value"]) + if any(choice.get("text") for choice in payload.get("choices", [])): + text_times.append(timestamp) + if len(done_times) != 1: + raise ValueError("Expected one complete SSE DONE marker") + + if abs((done_times[0] - origin) - (end - start)) > 1_000_000: + raise ValueError("HTTP completion and SSE DONE elapsed times disagree") + completion = (end - start) / 1e6 + first = (text_times[0] - origin) / 1e6 if text_times else None + last = (text_times[-1] - origin) / 1e6 if text_times else None + return dict( + http_completion_ms=completion, + first_visible_text_ms=first, + last_visible_text_ms=last, + done_after_last_visible_text_ms=(done_times[0] - text_times[-1]) / 1e6 if text_times else None, + ) + + +def summarize_streaming_times(records): + rows = [streaming_completion_times(record) for record in records] + if not rows: + raise ValueError("No streaming records to summarize") + content = [row for row in rows if row["last_visible_text_ms"] is not None] + return dict( + requests=len(rows), + empty_text_responses=len(rows) - len(content), + http_completion_all_requests_ms=_stats([row["http_completion_ms"] for row in rows]), + http_completion_content_requests_ms=_stats([row["http_completion_ms"] for row in content]), + first_visible_text_ms=_stats([row["first_visible_text_ms"] for row in content]), + last_visible_text_ms=_stats([row["last_visible_text_ms"] for row in content]), + done_after_last_visible_text_ms=_stats([row["done_after_last_visible_text_ms"] for row in content]), + ) diff --git a/test/benchmark/service/token_stream.py b/test/benchmark/service/token_stream.py new file mode 100644 index 0000000000..fc63f66afd --- /dev/null +++ b/test/benchmark/service/token_stream.py @@ -0,0 +1,237 @@ +import codecs +from dataclasses import dataclass, field +import json +import math + + +class SSEDecoder: + def __init__(self): + self.decoder = codecs.getincrementaldecoder("utf-8")() + self.buffer = "" + self.data = [] + + def feed(self, chunk: bytes) -> list[str]: + self.buffer += self.decoder.decode(chunk) + events = [] + while "\n" in self.buffer: + line, self.buffer = self.buffer.split("\n", 1) + line = line.removesuffix("\r") + if not line: + if self.data: + events.append("\n".join(self.data)) + self.data.clear() + elif line.startswith("data:"): + self.data.append(line[5:].removeprefix(" ")) + return events + + def finish(self): + self.buffer += self.decoder.decode(b"", final=True) + if self.buffer.strip() or self.data: + raise ValueError("Truncated SSE event at EOF") + + +@dataclass +class TokenStream: + framework: str + prompt_ids: list[int] + output_length: int + started: float + token_ids: list[int] = field(default_factory=list, init=False) + arrivals: list[float] = field(default_factory=list, init=False) + finished: bool = field(default=False, init=False) + done: bool = field(default=False, init=False) + events: int = field(default=0, init=False) + last_received: float = field(default=0.0, init=False, repr=False) + cached_input_tokens: int | None = field(default=None, init=False) + + def __post_init__(self): + if self.framework not in ("lightllm", "vllm", "sglang") or self.output_length < 1: + raise ValueError("Expected a supported framework and positive output length") + if not self.prompt_ids or any(type(token) is not int or token < 0 for token in self.prompt_ids): + raise ValueError("Prompt must contain exact integer token IDs") + self.last_received = self.started + if not math.isfinite(self.started): + raise ValueError("Start timestamp must be finite") + + def record(self, payload: str, received: float): + if not math.isfinite(received) or received < self.last_received: + raise ValueError("Token arrival timestamps must be finite and monotonic") + self.last_received = received + if payload == "[DONE]": + if not self.finished: + raise ValueError("Stream ended before the terminal generation event") + self.done = True + return + if self.done: + raise ValueError("Received data after DONE") + body = json.loads(payload) + if body.get("error"): + raise ValueError(f"Service stream error: {body['error']}") + self.events += 1 + if self.framework == "lightllm": + token = body["token"] + if token["prompt_tokens"] != len(self.prompt_ids): + raise ValueError("Stream prompt count differs from actual request IDs") + if token["count_output_tokens"] != len(self.token_ids) + 1: + raise ValueError("Native stream skipped or duplicated an output position") + ids = [token["id"]] + self._record_cache_count(token.get("prompt_cache_len")) + reason = body.get("finish_reason") if body.get("finished") is True else None + else: + choices = body.get("choices", []) + if len(choices) > 1 or (choices and choices[0].get("index") != 0): + raise ValueError("Benchmark requires one completion choice") + choice = choices[0] if choices else {} + ids = choice.get("token_ids") + if ids is None: + ids = [] + if choice.get("text") and not ids: + raise ValueError("Nonempty text event omitted explicit token IDs") + prompt = choice.get("prompt_token_ids") + if prompt is not None and prompt != self.prompt_ids: + raise ValueError("Service changed the exact tokenized prompt") + reason = choice.get("finish_reason") + if not isinstance(ids, list) or any(type(token) is not int or token < 0 for token in ids): + raise ValueError("Stream must expose integer token IDs") + if ids and self.finished: + raise ValueError("Received output tokens after terminal generation") + self.token_ids.extend(ids) + self.arrivals.extend([received] * len(ids)) + if len(self.token_ids) > self.output_length: + raise ValueError("Service exceeded the requested output length") + if reason is not None: + if reason != "length" or len(self.token_ids) != self.output_length: + raise ValueError("Fixed-length stream did not finish at its requested token count") + self.finished = True + usage = body.get("usage") + if usage is not None: + self._record_cache_count((usage.get("prompt_tokens_details") or {}).get("cached_tokens")) + if usage is not None and ( + usage["prompt_tokens"] != len(self.prompt_ids) or usage["completion_tokens"] != len(self.token_ids) + ): + raise ValueError("Stream usage does not match observed exact token counts") + + def _record_cache_count(self, count): + if count is None: + return + if type(count) is not int or not 0 <= count <= len(self.prompt_ids): + raise ValueError("Cached input count must be an integer within the exact prompt") + if self.cached_input_tokens is not None and count != self.cached_input_tokens: + raise ValueError("Cached input count changed within one request") + self.cached_input_tokens = count + + def summarize(self, ended: float) -> dict: + if not self.finished or len(self.token_ids) != self.output_length: + raise ValueError("Incomplete generation cannot be included in benchmark results") + if self.framework != "lightllm" and not self.done: + raise ValueError("OpenAI-compatible stream omitted DONE") + if not math.isfinite(ended) or ended < self.last_received: + raise ValueError("Invalid request completion timestamp") + gaps = [right - left for left, right in zip(self.arrivals, self.arrivals[1:])] + return { + "input_tokens": len(self.prompt_ids), + "cached_input_tokens": self.cached_input_tokens, + "uncached_input_tokens": ( + len(self.prompt_ids) - self.cached_input_tokens if self.cached_input_tokens is not None else None + ), + "output_tokens": len(self.token_ids), + "events": self.events, + "ttft_s": self.arrivals[0] - self.started, + "tpot_s": sum(gaps) / len(gaps) if gaps else None, + "e2e_s": ended - self.started, + "receiver_itl_s": gaps, + "itl_definition": "client arrival gaps; IDs delivered in one event share a timestamp", + } + + +def request_stream( + base_url, + framework, + model, + prompt_ids, + output_length, + timeout, + response_path, + *, + disable_prompt_cache=False, + cache_salt=None, +): + import time + import requests + + if disable_prompt_cache and framework != "lightllm": + raise ValueError("Per-request prompt-cache disabling is only implemented for native requests") + if cache_salt is not None: + if framework not in ("vllm", "sglang"): + raise ValueError("Cache salts are only implemented for baseline requests") + if ( + not isinstance(cache_salt, str) + or not 1 <= len(cache_salt) <= 128 + or any(char in cache_salt for char in "@/\\\x00") + ): + raise ValueError("Cache salt must contain 1--128 characters without '@', '/', backslash or NUL") + if framework == "lightllm": + endpoint = "/generate_stream" + payload = { + "inputs": prompt_ids, + "parameters": { + "do_sample": False, + "ignore_eos": True, + "max_new_tokens": output_length, + "return_details": True, + }, + } + if disable_prompt_cache: + payload["parameters"]["disable_prompt_cache"] = True + elif framework in ("vllm", "sglang"): + endpoint = "/v1/completions" + payload = { + "model": model, + "prompt": prompt_ids, + "temperature": 0, + "seed": 3107, + "ignore_eos": True, + "max_tokens": output_length, + "stream": True, + "stream_options": {"include_usage": True}, + "return_token_ids": True, + } + if cache_salt is not None: + payload["cache_salt"] = cache_salt + else: + raise ValueError("Unsupported serving framework") + raw = [] + decoder = SSEDecoder() + try: + with requests.Session() as client: + client.trust_env = False + started = time.perf_counter() + stream = TokenStream(framework, prompt_ids, output_length, started) + with client.post(base_url.rstrip("/") + endpoint, json=payload, stream=True, timeout=timeout) as response: + if response.status_code != 200: + raw.append(response.content) + response.raise_for_status() + if not response.headers.get("content-type", "").startswith("text/event-stream"): + raw.append(response.content) + raise ValueError("Expected an SSE response") + if "chunked" not in response.headers.get("transfer-encoding", "").lower(): + raise ValueError("Token timing requires HTTP/1.1 chunked transfer") + for chunk in response.iter_content(chunk_size=None): + received = time.perf_counter() + raw.append(chunk) + for event in decoder.feed(chunk): + stream.record(event, received) + ended = time.perf_counter() + decoder.finish() + return { + "token_ids": stream.token_ids, + "metrics": stream.summarize(ended), + "timeline": { + "started": started, + "first_token": stream.arrivals[0], + "last_token": stream.arrivals[-1], + "ended": ended, + }, + } + finally: + response_path.write_bytes(b"".join(raw)) diff --git a/unit_tests/benchmark/qwen38/test_agentx_protocol.py b/unit_tests/benchmark/qwen38/test_agentx_protocol.py new file mode 100644 index 0000000000..4da19e1f81 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_agentx_protocol.py @@ -0,0 +1,223 @@ +import copy +import json + +import pytest + +from test.benchmark.qwen38.audit_agentx import audit, repeatability +from test.benchmark.qwen38.run_agentx import fixed_config, idle_window, reset_cache +from test.benchmark.qwen38.prepare_agentx import prompt_length +from test.benchmark.qwen38.prepare_agentx import frozen_sessions +from types import SimpleNamespace + + +def fixture(tmp_path, change=None): + payload = dict(model="model", messages=[dict(role="user", content="hello")], max_tokens=2) + inputs = dict(data=[dict(session_id="one", payloads=[payload])]) + row = dict( + metadata=dict( + benchmark_phase="profiling", + conversation_id="one", + turn_index=0, + request_start_ns=1_000_000_000, + request_end_ns=2_000_000_000, + ), + status=200, + payload=copy.deepcopy(payload), + responses=[ + dict( + packets=[ + dict( + name="data", + value=json.dumps( + dict( + choices=[dict(finish_reason="length")], usage=dict(prompt_tokens=3, completion_tokens=2) + ) + ), + ), + dict(name="data", value="[DONE]"), + ] + ) + ], + ) + summary = dict( + request_count=dict(avg=1), + usage_completion_tokens=dict(sum=2), + usage_prompt_tokens=dict(sum=3), + output_token_throughput=dict(avg=2), + request_throughput=dict(avg=1), + time_to_first_token=dict(p50=100, p95=120), + error_summary=[], + ) + if change: + change(row, summary) + (tmp_path / "raw_records").mkdir() + (tmp_path / "raw_records" / "worker.jsonl").write_text(json.dumps(row) + "\n") + (tmp_path / "profile_export_aiperf.json").write_text(json.dumps(summary)) + return inputs + + +def test_complete_fixed_work_is_independently_counted(tmp_path): + result = audit(tmp_path, fixture(tmp_path)) + assert result["requests"] == 1 + assert result["output_tokens"] == 2 + assert result["independent_output_tps"] == 2 + + +@pytest.mark.parametrize( + "failure", ["payload", "usage", "done", "error", "throughput", "request_throughput", "cancelled", "turn", "time"] +) +def test_audit_rejects_changed_work_or_miscounted_stream(tmp_path, failure): + def change(row, summary): + if failure == "payload": + row["payload"]["messages"][0]["content"] = "different" + elif failure == "usage": + packet = row["responses"][0]["packets"][0] + body = json.loads(packet["value"]) + body["usage"]["completion_tokens"] = 1 + packet["value"] = json.dumps(body) + elif failure == "done": + row["responses"][0]["packets"].pop() + elif failure == "error": + summary["error_summary"] = [dict(count=1)] + elif failure == "throughput": + summary["output_token_throughput"]["avg"] = 3 + elif failure == "request_throughput": + summary["request_throughput"]["avg"] = 2 + elif failure == "turn": + row["metadata"]["turn_index"] = 1 + elif failure == "time": + row["metadata"]["request_end_ns"] = 0 + else: + row["metadata"]["was_cancelled"] = True + + inputs = fixture(tmp_path, change) + with pytest.raises(ValueError): + audit(tmp_path, inputs) + + +def test_repeatability_rejects_large_variance_and_workload_mismatch(): + base = dict( + payloads_sha256="same", + prompt_counts_sha256="same-counts", + requests=100, + prompt_tokens=1000, + output_tokens=200, + output_tps=100, + request_tps=50, + ttft_p50_ms=10, + ) + assert repeatability([base, base, base])["status"] == "passed" + varied = dict(base, output_tps=190) + assert repeatability([base, varied, base])["status"] == "nonrepeatable" + with pytest.raises(ValueError, match="workload mismatch"): + repeatability([base, dict(base, prompt_tokens=1001), base]) + + +def test_fixed_work_has_no_time_cutoff_or_hidden_warmup(tmp_path): + config = fixed_config("http://localhost", "model", tmp_path, tmp_path / "inputs.json", 32, 32, tmp_path) + phases = config["benchmark"]["phases"] + assert len(phases) == 1 and phases[0]["sessions"] == 32 + assert "duration" not in phases[0] + assert "multiRun" not in config and "scenario" not in config["benchmark"] + + +@pytest.mark.parametrize("encoded", [[1, 2, 3], dict(input_ids=[1, 2, 3], attention_mask=[1, 1, 1])]) +def test_chat_template_length_handles_list_and_batch_encoding(encoded): + class Tokenizer: + def apply_chat_template(self, messages, add_generation_prompt): + return encoded + + assert prompt_length(Tokenizer(), []) == 3 + + +def test_freeze_reconstructs_deltas_and_resets_without_live_responses(): + def turn(text, reset=False): + return SimpleNamespace( + raw_messages=[dict(role="user", content=text)], + reset_context=reset, + max_tokens=5, + source_trace_id="trace", + source_outer_idx=0, + source_kind="weka_main", + ) + + conversation = SimpleNamespace(session_id="one", turns=[turn("first"), turn("second"), turn("new", True)]) + inputs, sources = frozen_sessions([conversation], "model", 1, 3) + payloads = inputs["data"][0]["payloads"] + assert [p["messages"] for p in payloads] == [ + [dict(role="user", content="first")], + [dict(role="user", content="first"), dict(role="user", content="second")], + [dict(role="user", content="new")], + ] + assert len(sources[0]["turns"]) == 3 + + +def test_audit_uses_merged_export_without_double_counting_worker_shards(tmp_path): + inputs = fixture(tmp_path) + row = (tmp_path / "raw_records" / "worker.jsonl").read_text() + (tmp_path / "profile_export_raw.jsonl").write_text(row) + assert audit(tmp_path, inputs)["requests"] == 1 + + +def test_audit_counts_completion_usage_including_reasoning(tmp_path): + inputs = fixture(tmp_path) + path = tmp_path / "profile_export_aiperf.json" + summary = json.loads(path.read_text()) + summary["output_token_count"] = dict(sum=0) + summary["reasoning_token_count"] = dict(sum=2) + path.write_text(json.dumps(summary)) + assert audit(tmp_path, inputs)["output_tokens"] == 2 + + +def test_audit_rejects_overlapping_turns_with_identical_payloads(tmp_path): + inputs = fixture(tmp_path) + inputs["data"][0]["payloads"] *= 2 + path = tmp_path / "raw_records" / "worker.jsonl" + first = json.loads(path.read_text()) + second = copy.deepcopy(first) + second["metadata"].update(turn_index=1, request_start_ns=1_500_000_000, request_end_ns=2_500_000_000) + path.write_text(json.dumps(first) + "\n" + json.dumps(second) + "\n") + with pytest.raises(ValueError, match="overlapped"): + audit(tmp_path, inputs) + + +def test_sglang_cache_reset_text(monkeypatch): + target = "test.benchmark.qwen38.run_agentx.http" + monkeypatch.setattr(target, lambda *args: b"Cache flushed.\nPlease check backend logs for more details.\n") + assert reset_cache("http://localhost", "sglang")["path"] == "/flush_cache" + monkeypatch.setattr(target, lambda *args: b"Cache not flushed because there are pending requests.") + with pytest.raises(ValueError, match="Unexpected SGLang"): + reset_cache("http://localhost", "sglang") + + +@pytest.mark.parametrize("failure", ["transient", "busy", "foreign"]) +def test_idle_gate(monkeypatch, failure): + target = "test.benchmark.qwen38.run_agentx" + now = [0] + calls = [0] + + def output(command, **kwargs): + if command[0] == "sudo": + return "PID\n123\n" + if "--query-compute-apps=pid" in command: + return "456\n" if failure == "foreign" else "123\n" + calls[0] += 1 + busy = failure == "busy" or calls[0] in (1, 3) + return "\n".join(f"{i}, {9 if busy else 0}, 54" for i in range(8)) + + monkeypatch.setattr(target + ".subprocess.check_output", output) + monkeypatch.setattr(target + ".time.monotonic", lambda: now[0]) + monkeypatch.setattr(target + ".time.sleep", lambda seconds: now.__setitem__(0, now[0] + seconds)) + history = [] + if failure == "foreign": + with pytest.raises(ValueError, match="outside"): + idle_window("container", history) + assert not history + elif failure == "busy": + with pytest.raises(TimeoutError, match="idle"): + idle_window("container", history, timeout=60) + assert len(history) == 6 + else: + samples = idle_window("container", history) + assert len(history) == 6 and len(samples) == 3 + assert all(util == 0 for sample in samples for _, util, _ in sample) diff --git a/unit_tests/benchmark/qwen38/test_aiperf_replay.py b/unit_tests/benchmark/qwen38/test_aiperf_replay.py new file mode 100644 index 0000000000..2217ea9732 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_aiperf_replay.py @@ -0,0 +1,139 @@ +import copy +import json + +import pytest + +from test.benchmark.qwen38.measure_aiperf import audit_records +from test.benchmark.qwen38.prepare_aiperf_trace import select_requests + + +def test_trace_selection_preserves_bursts_prefixes_and_lengths(): + rows = [ + dict(timestamp=0, input_length=4000, output_length=200, hash_ids=[1]), + dict(timestamp=10, input_length=20, output_length=3, hash_ids=[2]), + dict(timestamp=10, input_length=30, output_length=4, hash_ids=[2, 3]), + dict(timestamp=100, input_length=10, output_length=5, hash_ids=[4]), + ] + selected, count = select_requests(rows, 2, 4096) + assert count == 3 and selected == [(1, rows[1]), (2, rows[2])] + with pytest.raises(ValueError, match="Not enough"): + select_requests(rows, 4, 4096) + with pytest.raises(ValueError, match="nondecreasing"): + select_requests(list(reversed(rows)), 1, 4096) + + +def fixture_record(): + return dict( + metadata=dict( + benchmark_phase="profiling", + conversation_id="session_000000", + session_num=0, + request_start_ns=1000, + was_cancelled=False, + ), + status=200, + payload=dict(prompt=["hello"], max_tokens=2, temperature=0, ignore_eos=True), + responses=[ + dict( + packets=[ + dict( + name="data", + value=json.dumps( + dict( + choices=[dict(text="hello", finish_reason="length")], + usage=dict(prompt_tokens=1, completion_tokens=2), + ) + ), + ), + dict(name="data", value="[DONE]"), + ] + ) + ], + ) + + +def write_artifacts(path, records): + metrics = ( + "request_throughput", + "output_token_throughput", + "time_to_first_token", + "inter_token_latency", + "request_latency", + ) + summary = dict(error_summary=[], **{name: dict(avg=1) for name in metrics}) + summary.update(request_count=dict(avg=len(records)), output_token_count=dict(sum=2 * len(records))) + (path / "profile_export_aiperf.json").write_text(json.dumps(summary)) + (path / "profile_export_raw.jsonl").write_text("".join(json.dumps(row) + "\n" for row in records)) + + +def test_raw_audit_accepts_complete_stream_and_excludes_model_alias(tmp_path): + row = fixture_record() + write_artifacts(tmp_path, [row]) + trace = [dict(timestamp=0, output_length=2)] + result = audit_records(tmp_path, trace) + row["payload"]["model"] = "different-service-alias" + write_artifacts(tmp_path, [row]) + assert audit_records(tmp_path, trace)["request_manifest_sha256"] == result["request_manifest_sha256"] + + +@pytest.mark.parametrize("failure", ["missing", "duplicate", "cancelled", "http", "no_done", "count"]) +def test_raw_audit_rejects_invalid_work(tmp_path, failure): + row = fixture_record() + rows = [row] + if failure == "missing": + rows = [] + elif failure == "duplicate": + rows.append(copy.deepcopy(row)) + elif failure == "cancelled": + row["metadata"]["was_cancelled"] = True + elif failure == "http": + row["status"] = 500 + elif failure == "no_done": + row["responses"][0]["packets"].pop() + elif failure == "count": + packet = row["responses"][0]["packets"][0] + body = json.loads(packet["value"]) + body["usage"]["completion_tokens"] = 1 + packet["value"] = json.dumps(body) + write_artifacts(tmp_path, rows) + with pytest.raises(ValueError): + audit_records(tmp_path, [dict(timestamp=0, output_length=2)]) + + +@pytest.mark.parametrize("error_type", ["InvalidInferenceResultError", "TimeoutError"]) +def test_complete_empty_text_stays_an_explicit_error(tmp_path, error_type): + empty, valid = fixture_record(), fixture_record() + valid["metadata"].update(session_num=1, conversation_id="session_000001", request_start_ns=2000) + packet = empty["responses"][0]["packets"][0] + body = json.loads(packet["value"]) + body["choices"][0]["text"] = "" + packet["value"] = json.dumps(body) + write_artifacts(tmp_path, [empty, valid]) + path = tmp_path / "profile_export_aiperf.json" + summary = json.loads(path.read_text()) + summary.update( + request_count=dict(avg=1), + output_token_count=dict(sum=2), + error_summary=[dict(count=1, error_details=dict(type=error_type, message="No responses with actual content"))], + ) + path.write_text(json.dumps(summary)) + trace = [dict(timestamp=0, output_length=2), dict(timestamp=1, output_length=2)] + if error_type == "TimeoutError": + with pytest.raises(ValueError, match="other request error"): + audit_records(tmp_path, trace) + else: + result = audit_records(tmp_path, trace) + assert result["error_rate"] == 0.5 and result["empty_content_sessions"] == [0] + assert result["output_tokens"] == 4 and result["successful_output_tokens"] == 2 + + +def test_concurrent_dispatch_order_does_not_reorder_source_requests(tmp_path): + first, second = fixture_record(), fixture_record() + second["metadata"].update(session_num=1, conversation_id="session_000001", request_start_ns=2000) + second["payload"]["prompt"] = ["different"] + trace = [dict(timestamp=0, output_length=2), dict(timestamp=1, output_length=2)] + write_artifacts(tmp_path, [first, second]) + reference = audit_records(tmp_path, trace) + first["metadata"]["session_num"], second["metadata"]["session_num"] = 1, 0 + write_artifacts(tmp_path, [second, first]) + assert audit_records(tmp_path, trace) == reference diff --git a/unit_tests/benchmark/qwen38/test_check_cold_streams.py b/unit_tests/benchmark/qwen38/test_check_cold_streams.py new file mode 100644 index 0000000000..abd432cfc1 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_check_cold_streams.py @@ -0,0 +1,61 @@ +from copy import deepcopy + +import pytest + +from test.benchmark.qwen38.check_cold_streams import workloads_from_reference +from test.benchmark.qwen38.check_generation import REVISION + + +def reference(): + return dict( + status="recorded", + revision=REVISION, + output_length=2, + cases=dict(english=[1, 2], long_3=[1, 2, 3], repeat_long_3=[1, 2, 3]), + results={name: dict(token_ids=[7, 8]) for name in ("english", "long_3", "repeat_long_3")}, + ) + + +def test_qualified_workloads_keep_order_and_exclude_repeat_controls(): + evidence = reference() + workloads = workloads_from_reference(evidence, evidence["cases"]) + assert [(work.name, work.prompt_ids, work.expected_ids) for work in workloads] == [ + ("english", (1, 2), (7, 8)), + ("long_3", (1, 2, 3), (7, 8)), + ] + + +@pytest.mark.parametrize( + "field,value", [("status", "failed"), ("revision", "other"), ("output_length", True), ("output_length", 0)] +) +def test_failed_or_unpinned_reference_cannot_authorize_timing(field, value): + evidence = reference() + expected = deepcopy(evidence["cases"]) + evidence[field] = value + with pytest.raises(ValueError): + workloads_from_reference(evidence, expected) + + +@pytest.mark.parametrize( + "change", ["prompt", "missing_case", "extra_case", "truncated", "invalid_id", "repeat_prompt", "repeat_output"] +) +def test_incomplete_or_edited_evidence_cannot_authorize_timing(change): + evidence = reference() + expected = deepcopy(evidence["cases"]) + if change == "prompt": + evidence["cases"]["english"] = [99, 2] + elif change == "missing_case": + del evidence["results"]["long_3"] + elif change == "extra_case": + evidence["results"]["other"] = dict(token_ids=[7, 8]) + elif change == "truncated": + evidence["results"]["long_3"]["token_ids"] = [7] + elif change == "invalid_id": + evidence["results"]["long_3"]["token_ids"] = [7, True] + elif change == "repeat_prompt": + expected["repeat_long_3"] = [2, 2, 3] + evidence["cases"]["repeat_long_3"] = list(expected["repeat_long_3"]) + elif change == "repeat_output": + evidence["results"]["repeat_long_3"]["token_ids"] = [7, 9] + with pytest.raises(ValueError): + workloads_from_reference(evidence, expected) diff --git a/unit_tests/benchmark/qwen38/test_check_cpu_cache.py b/unit_tests/benchmark/qwen38/test_check_cpu_cache.py new file mode 100644 index 0000000000..cff922e327 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_check_cpu_cache.py @@ -0,0 +1,66 @@ +import copy +import importlib +from pathlib import Path + +import pytest + + +@pytest.fixture +def checker(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[3] / "test/benchmark/qwen38")) + return importlib.import_module("check_cpu_cache") + + +def events(): + result = [] + for rank in range(4): + common = dict(rank=rank, first_token=1000, request_id=80) + result.extend( + [ + dict(common, event="cpu_offload_queued", monotonic=1, pages=10), + dict( + common, + event="cpu_extra_load_queued", + monotonic=4.5, + lexical_layers=[1], + slot_fields=["qsa_raw_keys", "qsa_pooled_keys", "qsa_key_cos", "qsa_key_sin"], + ), + dict(common, event="cpu_load_queued", monotonic=5, tokens=2368, kv_length=2368, pinned=True), + ] + ) + return result + + +def test_complete_target_cpu_operations_required_on_every_rank(checker): + proof = checker.validate_cpu_restore(events(), 4, 1000, 0.5, (4, 6)) + assert proof == [dict(rank=rank, offloads=1, restores=1) for rank in range(4)] + + +@pytest.mark.parametrize("kind", ["cpu_offload_queued", "cpu_load_queued", "cpu_extra_load_queued"]) +def test_partial_transfer_evidence_cannot_pass(checker, kind): + incomplete = [event for event in events() if not (event["rank"] == 3 and event["event"] == kind)] + with pytest.raises(ValueError, match="TP rank 3"): + checker.validate_cpu_restore(incomplete, 4, 1000, 0.5, (4, 6)) + + +@pytest.mark.parametrize( + "kind, field, value", + [ + ("cpu_load_queued", "pinned", False), + ("cpu_load_queued", "tokens", 0), + ("cpu_load_queued", "kv_length", 64), + ("cpu_load_queued", "monotonic", 3), + ("cpu_load_queued", "request_id", None), + ("cpu_extra_load_queued", "request_id", 81), + ("cpu_extra_load_queued", "slot_fields", ["qsa_raw_keys"]), + ("cpu_extra_load_queued", "lexical_layers", []), + ("cpu_offload_queued", "first_token", 1001), + ], +) +def test_stale_wrong_request_or_incomplete_model_fields_cannot_pass(checker, kind, field, value): + altered = copy.deepcopy(events()) + for event in altered: + if event["rank"] == 0 and event["event"] == kind: + event[field] = value + with pytest.raises(ValueError, match="TP rank 0"): + checker.validate_cpu_restore(altered, 4, 1000, 0.5, (4, 6)) diff --git a/unit_tests/benchmark/qwen38/test_check_generation.py b/unit_tests/benchmark/qwen38/test_check_generation.py new file mode 100644 index 0000000000..7e9f38ce5c --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_check_generation.py @@ -0,0 +1,95 @@ +import hashlib +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from test.benchmark.qwen38 import check_generation + + +@pytest.fixture +def tokenizer_root(tmp_path, monkeypatch): + payload = b"pinned tokenizer fixture" + (tmp_path / "tokenizer.json").write_bytes(payload) + monkeypatch.setattr(check_generation, "TOKENIZER_SHA256", hashlib.sha256(payload).hexdigest()) + + class Tokenizer: + def encode(self, prompt, **kwargs): + return SimpleNamespace(ids=[7, 8, 9]) + + monkeypatch.setattr(check_generation.Tokenizer, "from_file", lambda path: Tokenizer()) + return tmp_path + + +def test_default_workload_preserves_prior_long_cases(tokenizer_root): + cases = check_generation.cases_for(tokenizer_root) + assert list(cases) == [ + "english", + "chinese", + "code_unicode", + "long_2049", + "repeat_long_2049", + "long_4063", + "repeat_long_4063", + ] + assert len(cases["long_2049"]) == 2049 + assert len(cases["long_4063"]) == 4063 + + +def test_custom_lengths_keep_exact_repeated_inputs_without_aliasing(tokenizer_root): + cases = check_generation.cases_for(tokenizer_root, (16, 31)) + assert len(cases["long_16"]) == 16 and len(cases["long_31"]) == 31 + assert cases["long_16"] == cases["repeat_long_16"] + assert cases["long_31"] == cases["repeat_long_31"] + cases["repeat_long_16"][0] = 100 + assert cases["long_16"][0] == 7 + assert "long_2049" not in cases + + +@pytest.mark.parametrize("lengths", [(), (0,), (-1,), (True,), (2.0,), (16, 16)]) +def test_invalid_lengths_fail_before_loading_tokenizer(lengths): + with pytest.raises(ValueError, match="Long prompt lengths"): + check_generation.cases_for(Path("missing-tokenizer"), lengths) + + +def test_custom_lengths_still_require_pinned_tokenizer(tokenizer_root): + (tokenizer_root / "tokenizer.json").write_bytes(b"changed") + with pytest.raises(ValueError, match="pinned checkpoint"): + check_generation.cases_for(tokenizer_root, (31,)) + + +@pytest.mark.parametrize("framework", ["lightllm", "vllm", "sglang"]) +@pytest.mark.parametrize("cached", [0, 1, None, False]) +def test_cold_generation_isolates_requests_and_requires_explicit_zero_cache(framework, cached, tmp_path): + body = ( + dict(prompt_tokens=2, tokens=[dict(id=7, prompt_cache_len=cached), dict(id=8, prompt_cache_len=cached)]) + if framework == "lightllm" + else dict( + choices=[dict(token_ids=[7, 8])], + usage=dict(prompt_tokens=2, prompt_tokens_details=dict(cached_tokens=cached)), + ) + ) + sent = [] + serialized = json.dumps(body) + + def post(url, json, timeout): + sent.append(json) + return SimpleNamespace(text=serialized, json=lambda: body, raise_for_status=lambda: None) + + client = SimpleNamespace(post=post) + for index in range(2): + raw = tmp_path / f"response-{index}.json" + arguments = (client, "http://service", framework, "model", [1, 2], 2, 10) + if type(cached) is int and cached == 0: + result = check_generation.generate(*arguments, response_path=raw, cold_cache=True) + assert result["token_ids"] == [7, 8] + else: + with pytest.raises(ValueError, match="zero cached"): + check_generation.generate(*arguments, response_path=raw, cold_cache=True) + assert json.loads(raw.read_text()) == body + if framework == "lightllm": + assert all(payload["parameters"]["disable_prompt_cache"] is True for payload in sent) + else: + assert all(isinstance(payload["cache_salt"], str) and payload["cache_salt"] for payload in sent) + assert sent[0]["cache_salt"] != sent[1]["cache_salt"] diff --git a/unit_tests/benchmark/qwen38/test_check_math_quality.py b/unit_tests/benchmark/qwen38/test_check_math_quality.py new file mode 100644 index 0000000000..5cb9681f71 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_check_math_quality.py @@ -0,0 +1,133 @@ +from copy import deepcopy +import hashlib +import json +import sys +from types import SimpleNamespace + +import pytest + +from test.benchmark.qwen38 import check_math_quality as check +from test.benchmark.qwen38.math_quality import MathExample + + +@pytest.fixture +def workload(tmp_path, monkeypatch): + files = {} + for name in ("tokenizer.json", "tokenizer_config.json", "chat_template.jinja"): + payload = name.encode() + (tmp_path / name).write_bytes(payload) + files[name] = hashlib.sha256(payload).hexdigest() + monkeypatch.setattr(check, "TOKENIZER_SHA256", files["tokenizer.json"]) + examples = (MathExample(0, "What is two plus two?", "4"),) + monkeypatch.setattr(check, "load_examples", lambda path: examples) + + class Tokenizer: + def apply_chat_template(self, *args, **kwargs): + return [1, 2] + + tokenizer = Tokenizer() + manifest = dict( + status="prepared", + model_revision=check.REVISION, + dataset=dict(revision=check.DATASET_REVISION, sha256=check.DATASET_SHA256, split="test", records=1), + tokenizer_files=files, + decoding=dict(temperature=0, ignore_eos=False, enable_thinking=False), + sample_count=1, + sample_seed=38, + max_output_tokens=3, + context_limit=5, + cases=check.prepare_cases(examples, tokenizer, 3, 5), + ) + return manifest, tmp_path, tokenizer + + +def test_verifies_reconstructed_prompts_and_gold_labels(workload): + manifest, root, tokenizer = workload + assert check.verified_workload(manifest, root, root / "data", tokenizer) == manifest["cases"] + for field, changed in [("expected_number", "999"), ("prompt_ids", [1, 3]), ("dataset_index", 1)]: + forged = deepcopy(manifest) + forged["cases"][0][field] = changed + with pytest.raises(ValueError, match="prompts or gold"): + check.verified_workload(forged, root, root / "data", tokenizer) + + +def test_rejects_changed_tokenizer_and_wrong_protocol(workload): + manifest, root, tokenizer = workload + forged = deepcopy(manifest) + forged["decoding"]["ignore_eos"] = True + with pytest.raises(ValueError, match="greedy non-thinking"): + check.verified_workload(forged, root, root / "data", tokenizer) + (root / "chat_template.jinja").write_bytes(b"changed") + with pytest.raises(ValueError, match="Tokenizer files"): + check.verified_workload(manifest, root, root / "data", tokenizer) + + +@pytest.mark.parametrize("bad_count", [False, True]) +def test_client_keeps_wrong_and_truncated_answers_in_denominator(tmp_path, monkeypatch, bad_count): + metadata = dict( + framework="lightllm", version="fixture", command=["fixture"], hardware="fixture", configuration={"tp": 1} + ) + (tmp_path / "service.json").write_text(json.dumps(metadata)) + (tmp_path / "manifest.json").write_text("{}") + (tmp_path / "generation_config.json").write_text(json.dumps(dict(eos_token_id=[99]))) + cases = [dict(dataset_index=i, prompt_ids=[1, 2], max_output_tokens=3, expected_number="4") for i in range(3)] + + class Tokenizer: + eos_token_id = 99 + + def decode(self, ids, **kwargs): + return {11: "#### 4", 12: "#### 5", 13: "#### 4"}[ids[0]] + + class Session: + index = 0 + + def post(self, url, *, json, timeout): + assert "expected_number" not in json and json["inputs"] == [1, 2] + ids, reason = [([11, 99], "stop"), ([12, 99], "stop"), ([13, 14, 15], "length")][self.index] + self.index += 1 + body = dict( + tokens=[dict(id=token) for token in ids], + prompt_tokens=2, + count_output_tokens=1 if bad_count else len(ids), + finish_reason=reason, + ) + return SimpleNamespace(text="fixture", raise_for_status=lambda: None, json=lambda: body) + + def close(self): + pass + + monkeypatch.setattr(check, "get_tokenizer", lambda path: SimpleNamespace(tokenizer=Tokenizer())) + monkeypatch.setattr(check, "verified_checkpoint", lambda path: None) + monkeypatch.setattr(check, "verified_workload", lambda *args: cases) + monkeypatch.setattr(check.requests, "Session", Session) + monkeypatch.setenv("EXP_ACTIVE", "1") + monkeypatch.setenv("EXP_RUN_DIR", str(tmp_path)) + monkeypatch.setattr( + sys, + "argv", + [ + "check_math_quality.py", + "--framework", + "lightllm", + "--base-url", + "http://fixture", + "--model-directory", + str(tmp_path), + "--service-metadata", + str(tmp_path / "service.json"), + "--manifest", + str(tmp_path / "manifest.json"), + "--dataset", + str(tmp_path / "data"), + ], + ) + if bad_count: + with pytest.raises(ValueError, match="counts"): + check.main() + report = json.loads((tmp_path / "math-quality.json").read_text()) + assert report["status"] == "failed" and "accuracy" not in report + else: + check.main() + report = json.loads((tmp_path / "math-quality.json").read_text()) + assert report["status"] == "recorded" + assert (report["scored"], report["correct"], report["truncated"], report["accuracy"]) == (3, 1, 1, 1 / 3) diff --git a/unit_tests/benchmark/qwen38/test_check_multimodal.py b/unit_tests/benchmark/qwen38/test_check_multimodal.py new file mode 100644 index 0000000000..928a72549c --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_check_multimodal.py @@ -0,0 +1,98 @@ +import importlib +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +@pytest.fixture +def checker(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[3] / "test/benchmark/qwen38")) + return importlib.import_module("check_multimodal") + + +def fixture(): + case = dict(prompt_tokens=90, image_tokens=64, max_new_tokens=64, ignore_eos=False, expected_words=["red"]) + body = dict( + prompt_tokens=90, + count_output_tokens=2, + finish_reason="stop", + tokens=[dict(id=7), dict(id=248046)], + input_usage=dict(input_text_tokens=26, input_image_tokens=64, input_audio_tokens=0), + ) + tokenizer = SimpleNamespace(decode=lambda ids, **_: " ".join({7: "Red", 8: "blue"}[value] for value in ids)) + return case, body, tokenizer + + +def test_image_count_eos_and_semantic_validation(checker): + case, body, tokenizer = fixture() + result = checker.validate_image_result(body, case, [248044, 248046], tokenizer) + assert result["token_ids"] == [7, 248046] and result["answer"] == "Red" + + +@pytest.mark.parametrize( + "field,value", [("input_text_tokens", 27), ("input_image_tokens", 63), ("input_audio_tokens", 1)] +) +def test_incorrect_modality_counts_cannot_pass(checker, field, value): + case, body, tokenizer = fixture() + body["input_usage"][field] = value + with pytest.raises(ValueError, match="accounting"): + checker.validate_image_result(body, case, [248046], tokenizer) + + +def test_wrong_color_and_swapped_image_order_cannot_pass(checker): + case, body, tokenizer = fixture() + body["tokens"][0]["id"] = 8 + with pytest.raises(ValueError, match="color/order"): + checker.validate_image_result(body, case, [248046], tokenizer) + case["expected_words"] = ["red", "blue"] + body["tokens"] = [dict(id=8), dict(id=7), dict(id=248046)] + body["count_output_tokens"] = 3 + with pytest.raises(ValueError, match="color/order"): + checker.validate_image_result(body, case, [248046], tokenizer) + + +def test_pressure_request_has_explicit_fixed_length_contract(checker): + case, body, tokenizer = fixture() + case.update(expected_words=[], max_new_tokens=1, ignore_eos=True) + body.update(tokens=[dict(id=7)], count_output_tokens=1, finish_reason="length") + assert checker.validate_image_result(body, case, [248046], tokenizer)["token_ids"] == [7] + + +def cache_events(): + first = dict(id=1, token_id=100000000, token_num=64, md5="digest_0") + last = dict(id=2, token_id=100000384, token_num=64, md5="digest_0") + return [ + dict(event="alloc", records=[first]), + dict(event="embedded", records=[first]), + dict(event="evicted", record=first), + dict(event="alloc", records=[last]), + dict(event="embedded", records=[last]), + ] + + +def test_reallocation_requires_new_virtual_ids_eviction_and_embedding_completion(checker): + proof = checker.verify_reallocation(cache_events(), "digest") + assert proof["original_evicted"] and proof["both_embedded"] + assert proof["original"]["token_id"] != proof["reallocated"]["token_id"] + + +@pytest.mark.parametrize("missing", ["alloc", "evicted", "embedded"]) +def test_incomplete_cache_evidence_cannot_pass(checker, missing): + events = [item for item in cache_events() if item["event"] != missing] + with pytest.raises(ValueError): + checker.verify_reallocation(events, "digest") + + +def test_cache_reallocation_must_change_virtual_range(checker): + events = cache_events() + events[3]["records"][0]["token_id"] = events[0]["records"][0]["token_id"] + with pytest.raises(ValueError, match="new virtual IDs"): + checker.verify_reallocation(events, "digest") + + +def test_fixture_hashes_distinguish_cache_pressure_images(checker, tmp_path): + payload, images = checker.image_fixtures(tmp_path, 5) + assert len(payload) == len(images) == 8 + assert len({image["sha256"] for image in images.values()}) == 8 + assert images["rect_blue"]["width"] == 384 and images["rect_blue"]["height"] == 128 diff --git a/unit_tests/benchmark/qwen38/test_check_prefix_eviction.py b/unit_tests/benchmark/qwen38/test_check_prefix_eviction.py new file mode 100644 index 0000000000..e5b4e80f53 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_check_prefix_eviction.py @@ -0,0 +1,86 @@ +import copy +import importlib +from pathlib import Path + +import pytest + + +@pytest.fixture +def checker(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[3] / "test/benchmark/qwen38")) + return importlib.import_module("check_prefix_eviction") + + +def events(): + result = [] + for rank in range(4): + result.extend( + [ + dict(event="state_freed", rank=rank, monotonic=2.0, pool=10, indexes=[3]), + dict(event="node_evicted", rank=rank, monotonic=2.1, first_token=1000, big_pool=10, big_state=3), + dict(event="state_freed", rank=rank, monotonic=3.0, pool=20, indexes=[4]), + dict( + event="small_state_evicted", + rank=rank, + monotonic=3.1, + first_token=1000, + small_pool=20, + small_state=4, + ), + ] + ) + return result + + +def test_actual_target_big_and_small_state_releases_required_on_every_rank(checker): + proof = checker.validate_eviction(events(), 4, 1000, 1, 4) + assert {item["rank"] for item in proof} == {0, 1, 2, 3} + + +@pytest.mark.parametrize("missing", ["node_evicted", "small_state_evicted", "state_freed"]) +def test_partial_eviction_evidence_cannot_pass(checker, missing): + incomplete = [event for event in events() if event["event"] != missing] + with pytest.raises(ValueError, match="No observed target"): + checker.validate_eviction(incomplete, 4, 1000, 1, 4) + + +def test_one_rank_trace_cannot_prove_four_rank_eviction(checker): + with pytest.raises(ValueError, match="rank 1"): + checker.validate_eviction([event for event in events() if event["rank"] == 0], 4, 1000, 1, 4) + + +@pytest.mark.parametrize("change", ["other_target", "other_pool", "other_index", "late_free", "outside_interval"]) +def test_unrelated_or_out_of_order_events_cannot_prove_target_eviction(checker, change): + changed = copy.deepcopy(events()) + for event in changed: + if change == "other_target" and "first_token" in event: + event["first_token"] = 1001 + elif change == "other_pool" and event["event"] == "state_freed": + event["pool"] += 50 + elif change == "other_index" and event["event"] == "state_freed": + event["indexes"] = [900] + elif change == "late_free" and event["event"] == "state_freed": + event["monotonic"] = 3.9 + elif change == "outside_interval": + event["monotonic"] += 10 + with pytest.raises(ValueError, match="No observed target"): + checker.validate_eviction(changed, 4, 1000, 1, 4) + + +@pytest.mark.parametrize("value", [-1, True, 2305, 4.0]) +def test_invalid_cache_hit_accounting_rejected(checker, value): + result = dict(response=dict(prompt_tokens=2305, tokens=[dict(prompt_cache_len=value)])) + with pytest.raises(ValueError, match="cache hit"): + checker.cache_hit(result) + + +def test_pressure_prefixes_diverge_before_first_hash_page(checker, monkeypatch): + monkeypatch.setattr(checker, "cases_for", lambda _: {"long_4063": [9] * 4063}) + cases = checker.eviction_cases(Path("unused")) + assert len({tuple(case["inputs"][:64]) for case in cases.values()}) == 33 + assert (len(cases["target"]["inputs"]) - 1) % 256 == 64 + assert all( + len(case["inputs"]) == (65 if name.startswith("small_pressure_") else 2369) + and len(case["inputs"]) + case["max_new_tokens"] <= 4096 + for name, case in cases.items() + ) diff --git a/unit_tests/benchmark/qwen38/test_check_scheduler.py b/unit_tests/benchmark/qwen38/test_check_scheduler.py new file mode 100644 index 0000000000..3ff7319ec7 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_check_scheduler.py @@ -0,0 +1,121 @@ +import importlib +import json +from pathlib import Path +import threading + +import pytest +import requests + + +@pytest.fixture +def client_module(monkeypatch): + root = Path(__file__).resolve().parents[3] + monkeypatch.syspath_prepend(str(root / "test/benchmark/qwen38")) + return importlib.import_module("check_scheduler") + + +def response_for(monkeypatch, module, body, http_error=False): + class Response: + text = json.dumps(body) + + def raise_for_status(self): + if http_error: + raise requests.HTTPError("503") + + def json(self): + return body + + class Session: + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def post(self, url, json, timeout): + assert json["parameters"]["do_sample"] is False + assert json["parameters"]["return_details"] is True + assert json["parameters"]["print_eos_token"] is True + return Response() + + monkeypatch.setattr(module.requests, "Session", Session) + + +@pytest.mark.parametrize( + "ids,ignore_eos,reason", [([3, 4], True, "length"), ([248046], False, "stop"), ([3, 248044], False, "stop")] +) +def test_valid_termination_preserves_raw_response(client_module, monkeypatch, tmp_path, ids, ignore_eos, reason): + body = { + "tokens": [{"id": token} for token in ids], + "prompt_tokens": 3, + "count_output_tokens": len(ids), + "finish_reason": reason, + } + response_for(monkeypatch, client_module, body) + case = {"inputs": [7, 8, 9], "max_new_tokens": 2, "ignore_eos": ignore_eos} + path = tmp_path / "response.txt" + result = client_module.generate("http://localhost", case, [248044, 248046], 10, path) + assert result["token_ids"] == ids and json.loads(path.read_text()) == body + + +@pytest.mark.parametrize( + "ids,ignore_eos,reason,prompt,count", + [ + ([], True, "length", 3, 0), + ([True, 4], True, "length", 3, 2), + ([-1, 4], True, "length", 3, 2), + ([3, 4, 5], True, "length", 3, 3), + ([3], True, "length", 3, 1), + ([3, 4], True, "stop", 3, 2), + ([3, 4], True, "length", 2, 2), + ([3, 4], True, "length", 3, 1), + ([3], False, "length", 3, 1), + ([248046, 3], False, "stop", 3, 2), + ([3, 248046], False, "length", 3, 2), + ], +) +def test_invalid_response_cannot_pass(client_module, monkeypatch, tmp_path, ids, ignore_eos, reason, prompt, count): + body = { + "tokens": [{"id": token} for token in ids], + "prompt_tokens": prompt, + "count_output_tokens": count, + "finish_reason": reason, + } + response_for(monkeypatch, client_module, body) + case = {"inputs": [7, 8, 9], "max_new_tokens": 2, "ignore_eos": ignore_eos} + path = tmp_path / "response.txt" + with pytest.raises(ValueError): + client_module.generate("http://localhost", case, [248044, 248046], 10, path) + assert json.loads(path.read_text()) == body + + +def test_http_error_is_saved(client_module, monkeypatch, tmp_path): + body = {"error": "service unavailable"} + response_for(monkeypatch, client_module, body, http_error=True) + path = tmp_path / "response.txt" + with pytest.raises(requests.HTTPError): + client_module.generate("http://localhost", {"inputs": [7]}, [248046], 10, path) + assert json.loads(path.read_text()) == body + + +def test_monitor_records_server_batch_gauges(client_module, monkeypatch, tmp_path): + stop = threading.Event() + payload = '# a comment\nlightllm_num_running_reqs{model_name="test"} 4.0\nlightllm_batch_current_size 3.0\n' + + class Session: + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def get(self, *args, **kwargs): + stop.set() + return type("Response", (), {"text": payload, "raise_for_status": lambda self: None})() + + monkeypatch.setattr(client_module.requests, "Session", Session) + summary = {"maximum_router_batch": 0, "samples": 0} + path = tmp_path / "metrics.txt" + client_module.monitor("http://localhost", stop, summary, path) + assert summary == {"maximum_router_batch": 4.0, "samples": 1} + assert path.read_text() == payload diff --git a/unit_tests/benchmark/qwen38/test_generation_observations.py b/unit_tests/benchmark/qwen38/test_generation_observations.py new file mode 100644 index 0000000000..ca27acc960 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_generation_observations.py @@ -0,0 +1,62 @@ +import copy + +import pytest + +from test.benchmark.qwen38.check_cold_streams import workloads_from_reference +from test.benchmark.qwen38.check_generation import REVISION +from test.benchmark.qwen38.generation_observations import workloads_from_observations + + +def observations(): + cases = {"long_16": [1, 2], "repeat_long_16": [1, 2]} + results = {} + for name, ids in [("long_16", [7, 8]), ("repeat_long_16", [7, 9])]: + results[name] = dict( + token_ids=ids, + response=dict( + choices=[dict(token_ids=ids)], + usage=dict(prompt_tokens=2, completion_tokens=2, prompt_tokens_details=dict(cached_tokens=0)), + ), + ) + return dict( + status="observations_recorded", + service=dict(framework="vllm"), + revision=REVISION, + cases=cases, + reference=None, + cache_policy="cold", + results=results, + output_length=2, + repeat_consistency=dict(status="failed", mismatches=[16]), + ) + + +def test_observations_retain_repeat_failure_and_cannot_qualify_a_reference(): + report = observations() + before = copy.deepcopy(report) + workloads, quality = workloads_from_observations(report, report["cases"]) + assert report == before + assert len(workloads) == 1 and workloads[0].expected_ids == (7, 8) + assert quality == dict(status="failed", mismatches=[16]) + with pytest.raises(ValueError, match="qualified"): + workloads_from_reference(report, report["cases"]) + + +@pytest.mark.parametrize("change", ["native", "cached", "count", "edited_ids", "false_pass", "incomplete"]) +def test_observations_reject_invalid_protocol_or_hidden_quality_failure(change): + report = observations() + result = report["results"]["long_16"] + if change == "native": + report["service"]["framework"] = "lightllm" + elif change == "cached": + result["response"]["usage"]["prompt_tokens_details"]["cached_tokens"] = 1 + elif change == "count": + result["response"]["usage"]["completion_tokens"] = 1 + elif change == "edited_ids": + result["response"]["choices"][0]["token_ids"] = [7, 10] + elif change == "false_pass": + report["repeat_consistency"] = dict(status="passed", mismatches=[]) + elif change == "incomplete": + del report["results"]["repeat_long_16"] + with pytest.raises(ValueError): + workloads_from_observations(report, report["cases"]) diff --git a/unit_tests/benchmark/qwen38/test_math_quality.py b/unit_tests/benchmark/qwen38/test_math_quality.py new file mode 100644 index 0000000000..1bf77ab80b --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_math_quality.py @@ -0,0 +1,145 @@ +import hashlib +import json +from copy import deepcopy + +import pytest + +from test.benchmark.qwen38 import math_quality as math + + +@pytest.mark.parametrize( + "value, expected", + [("1,234.00", "1234"), ("+001.2500", "1.25"), ("-0.00", "0"), ("-42", "-42"), ("0.0010", "0.001")], +) +def test_decimal_normalization_is_exact(value, expected): + assert math.normalize_number(value) == expected + large = "123456789012345678901234567890123456789" + assert math.normalize_number(large) == large + + +@pytest.mark.parametrize("value", ["1,23", "1.2.3", "1/2", "1e3", "NaN", "inf", "", None]) +def test_invalid_plain_numbers_are_not_guessed(value): + assert math.normalize_number(value) is None + + +def test_score_uses_final_numeric_answer_only(): + assert math.score_answer("Work: 10 - 4 = 6.\n#### 6.00\n", "6")["correct"] + assert not math.score_answer("Work: 6 + 1 = 7.\n#### 7", "6")["correct"] + with pytest.raises(ValueError, match="Expected answer"): + math.score_answer("#### 6", "unknown") + + +@pytest.mark.parametrize("text", ["The answer is 6", "#### 6\nMore work", "#### 7\n#### 6", "#### $6", "#### NaN"]) +def test_missing_ambiguous_or_nonfinal_answers_are_invalid(text): + assert math.extract_final_number(text) is None + + +def test_loader_pins_bytes_and_rejects_duplicate_questions(tmp_path, monkeypatch): + path = tmp_path / "test.jsonl" + rows = [ + dict(question="Question one", answer="A calculation.\n#### 4"), + dict(question="Question two", answer="#### 5"), + ] + payload = "\n".join(json.dumps(row) for row in rows).encode() + path.write_bytes(payload) + monkeypatch.setattr(math, "DATASET_SHA256", hashlib.sha256(payload).hexdigest()) + examples = math.load_examples(path) + assert [(row.index, row.expected) for row in examples] == [(0, "4"), (1, "5")] + path.write_bytes(payload + b"\n") + with pytest.raises(ValueError, match="pinned source"): + math.load_examples(path) + duplicate = "\n".join(json.dumps(rows[0]) for _ in range(2)).encode() + path.write_bytes(duplicate) + monkeypatch.setattr(math, "DATASET_SHA256", hashlib.sha256(duplicate).hexdigest()) + with pytest.raises(ValueError, match="distinct questions"): + math.load_examples(path) + + +def test_sampling_is_reproducible_and_without_replacement(): + examples = tuple(math.MathExample(i, f"Question {i}", "4") for i in range(20)) + selected = math.select_examples(examples, 8, 38) + assert selected == math.select_examples(examples, 8, 38) + assert len({row.index for row in selected}) == 8 + + +@pytest.mark.parametrize("limit, seed", [(0, 1), (2, 1), (True, 1), (1, True), (1, 1.0)]) +def test_invalid_selection_is_rejected(limit, seed): + with pytest.raises(ValueError): + math.select_examples((math.MathExample(0, "Question", "4"),), limit, seed) + + +def test_prompts_do_not_include_gold_solutions_and_preserve_token_budget(): + class Tokenizer: + messages = [] + + def apply_chat_template(self, messages, **kwargs): + self.messages.append(messages) + assert kwargs == dict(tokenize=True, add_generation_prompt=True, enable_thinking=False) + return [11, 12] + + tokenizer = Tokenizer() + first = math.prepare_cases([math.MathExample(0, "What is two plus two?", "4")], tokenizer, 3, 5) + second = math.prepare_cases([math.MathExample(0, "What is two plus two?", "999")], tokenizer, 3, 5) + assert tokenizer.messages[0] == tokenizer.messages[1] + assert first[0]["prompt_ids"] == second[0]["prompt_ids"] == [11, 12] + assert first[0]["expected_number"] == "4" and second[0]["expected_number"] == "999" + with pytest.raises(ValueError, match="serving context"): + math.prepare_cases([math.MathExample(0, "Question", "4")], tokenizer, 4, 5) + + +@pytest.mark.parametrize("ids", [[], [True], [-1], [1.0], "tokens"]) +def test_invalid_tokenizer_payloads_are_rejected(ids): + class Tokenizer: + def apply_chat_template(self, *args, **kwargs): + return ids + + with pytest.raises(ValueError, match="actual nonnegative"): + math.prepare_cases([math.MathExample(0, "Question", "4")], Tokenizer()) + + +def scored_report(): + return dict( + status="recorded", + manifest_sha256="a" * 64, + scored=2, + results=[ + dict( + dataset_index=1, token_ids=[11, 99], finish_reason="stop", truncated=False, correct=True, predicted="4" + ), + dict( + dataset_index=2, + token_ids=[12, 13], + finish_reason="length", + truncated=True, + correct=False, + predicted=None, + ), + ], + ) + + +def test_output_comparison_keeps_wrong_and_truncated_cases(): + reference = scored_report() + assert math.compare_outputs(reference, deepcopy(reference)) == dict(status="passed", cases=2, mismatches=[]) + for field, value in [("token_ids", [12, 14]), ("finish_reason", "stop"), ("truncated", False), ("correct", True)]: + actual = deepcopy(reference) + actual["results"][1][field] = value + assert math.compare_outputs(reference, actual) == dict(status="failed", cases=2, mismatches=[2]) + + +@pytest.mark.parametrize("field,value", [("status", "failed"), ("manifest_sha256", "b" * 64), ("scored", 1)]) +def test_output_comparison_rejects_incomplete_or_different_workloads(field, value): + reference = scored_report() + actual = deepcopy(reference) + actual[field] = value + with pytest.raises(ValueError): + math.compare_outputs(reference, actual) + + +def test_output_comparison_rejects_duplicate_or_different_cases(): + reference = scored_report() + actual = deepcopy(reference) + for index in (1, 3): + actual["results"][1]["dataset_index"] = index + with pytest.raises(ValueError): + math.compare_outputs(reference, actual) diff --git a/unit_tests/benchmark/qwen38/test_measure_fixed_work.py b/unit_tests/benchmark/qwen38/test_measure_fixed_work.py new file mode 100644 index 0000000000..95d6cc3ca9 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_measure_fixed_work.py @@ -0,0 +1,55 @@ +from copy import deepcopy + +import pytest + +from test.benchmark.qwen38.measure_fixed_work import summarize_fixed_work +from test.benchmark.service.stream_cohort import StreamingWorkload, run_cohort +from unit_tests.benchmark.service.test_stream_cohort import observation + + +def fixture(): + work = StreamingWorkload("prompt", [12, 13], [20, 21]) + result = observation(1, [2, 3], 4, ids=[22, 23]) + result["metrics"]["cached_input_tokens"] = 0 + return work, result + + +def test_id_mismatch_keeps_exact_gate_failed_but_measures_equal_work(tmp_path): + work, result = fixture() + cohort = run_cohort([work], 1, lambda *_: result, tmp_path) + before = deepcopy(cohort) + metrics = summarize_fixed_work(cohort, [work]) + assert metrics["output_tokens_per_s"] == pytest.approx(2 / 3) + assert cohort == before + assert cohort["status"] == "failed" and cohort["metrics"] is None + + +@pytest.mark.parametrize("cached", [None, True, -1, 1]) +def test_unknown_or_nonzero_cache_cannot_produce_timing(tmp_path, cached): + work, result = fixture() + result["metrics"]["cached_input_tokens"] = cached + cohort = run_cohort([work], 1, lambda *_: result, tmp_path) + with pytest.raises(ValueError): + summarize_fixed_work(cohort, [work]) + + +@pytest.mark.parametrize("fault", ["transport", "input_length", "output_length", "timeline", "missing"]) +def test_invalid_work_never_produces_timing(tmp_path, fault): + work, result = fixture() + if fault == "input_length": + result["metrics"]["input_tokens"] = 3 + elif fault == "output_length": + result["token_ids"] = [22] + elif fault == "timeline": + result["timeline"]["ended"] = 0 + + def request(*_): + if fault == "transport": + raise TimeoutError("timed out") + return result + + cohort = run_cohort([work], 1, request, tmp_path) + if fault == "missing": + cohort["records"] = [] + with pytest.raises(ValueError): + summarize_fixed_work(cohort, [work]) diff --git a/unit_tests/benchmark/qwen38/test_retrieval_quality.py b/unit_tests/benchmark/qwen38/test_retrieval_quality.py new file mode 100644 index 0000000000..9311c87d8f --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_retrieval_quality.py @@ -0,0 +1,62 @@ +import hashlib +import json + +import pytest + +from test.benchmark.qwen38 import retrieval_quality as quality + +EXPECTED = dict(ALPHA="00123", BETA="456", GAMMA="789") + + +def test_exact_digit_strings_and_all_records_are_required(): + assert quality.score_records(" ALPHA = 00123;\n BETA=456; GAMMA=789\n", EXPECTED)["correct"] + result = quality.score_records("ALPHA=123; BETA=456; GAMMA=789", EXPECTED) + assert not result["correct"] and result["records"] == dict(ALPHA=False, BETA=True, GAMMA=True) + + +@pytest.mark.parametrize( + "text", + [ + None, + "", + "ALPHA=00123; BETA=456", + "ALPHA=00123; BETA=456; GAMMA=789.\n", + "ALPHA=00123; BETA=456; GAMMA=789; ALPHA=123", + "Here: ALPHA=00123; BETA=456; GAMMA=789", + ], +) +def test_incomplete_or_ambiguous_answers_cannot_pass(text): + result = quality.score_records(text, EXPECTED) + assert result["predicted"] is None and not result["correct"] + + +def test_manifest_and_actual_tokenizer_are_pinned(tmp_path, monkeypatch): + (tmp_path / "config.json").write_text(json.dumps(dict(text_config=dict(max_position_embeddings=10)))) + files = {} + for name in ("tokenizer.json", "tokenizer_config.json", "chat_template.jinja"): + payload = name.encode() + (tmp_path / name).write_bytes(payload) + files[name] = hashlib.sha256(payload).hexdigest() + manifest = dict( + status="prepared", + model_revision=quality.REVISION, + context_limit=10, + decoding=dict(temperature=0, ignore_eos=False, enable_thinking=False), + tokenizer_files=files, + cases=[dict(prompt_ids=[1, 2], max_output_tokens=8, expected_records=EXPECTED)], + ) + path = tmp_path / "manifest.json" + path.write_text(json.dumps(manifest)) + monkeypatch.setattr( + quality, "MANIFEST_SHA256", hashlib.sha256(json.dumps(manifest, sort_keys=True).encode()).hexdigest() + ) + assert quality.load_workload(path, tmp_path) == manifest + manifest["cases"][0]["expected_records"]["ALPHA"] = "999" + path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match="pinned preparation"): + quality.load_workload(path, tmp_path) + manifest["cases"][0]["expected_records"]["ALPHA"] = "00123" + path.write_text(json.dumps(manifest)) + (tmp_path / "chat_template.jinja").write_bytes(b"changed") + with pytest.raises(ValueError, match="Tokenizer"): + quality.load_workload(path, tmp_path) diff --git a/unit_tests/benchmark/qwen38/test_score_reference.py b/unit_tests/benchmark/qwen38/test_score_reference.py new file mode 100644 index 0000000000..d265be3b48 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_score_reference.py @@ -0,0 +1,49 @@ +import importlib +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + + +@pytest.fixture +def scorer(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[3] / "test/benchmark/qwen38")) + return importlib.import_module("score_reference") + + +class RecordingModel: + def __init__(self): + self.calls = [] + + def __call__(self, **kwargs): + prior = kwargs["past_key_values"] or () + ids = tuple(kwargs["input_ids"][0].tolist()) + positions = kwargs["position_ids"][0].tolist() + assert positions == list(range(len(prior), len(prior) + len(ids))) + assert kwargs["attention_mask"].tolist() == [[1] * (len(prior) + len(ids))] + assert kwargs["use_cache"] and kwargs["logits_to_keep"] == 1 + self.calls.append(kwargs) + history = prior + ids + return SimpleNamespace(past_key_values=history, logits=torch.tensor([[sum(history)]])) + + +@pytest.mark.parametrize("chunk_size,expected_calls", [(None, 1), (1, 9), (4, 3), (16, 1)]) +def test_prefill_chunking_preserves_history_and_next_decode_position(scorer, chunk_size, expected_calls): + model = RecordingModel() + inputs = torch.arange(9).unsqueeze(0) + output = scorer.forward_reference(model, inputs, None, 9, chunk_size) + assert len(model.calls) == expected_calls + assert output.past_key_values == tuple(range(9)) + assert output.logits.item() == 36 + output = scorer.forward_reference(model, torch.tensor([[17]]), output.past_key_values, 10, chunk_size) + assert output.past_key_values == (*range(9), 17) + assert output.logits.item() == 53 + + +@pytest.mark.parametrize("count,length,chunk_size", [(0, 0, None), (2, 1, 1), (2, 2, 0)]) +def test_invalid_forward_lengths_are_rejected_before_call(scorer, count, length, chunk_size): + model = RecordingModel() + with pytest.raises(ValueError, match="valid lengths"): + scorer.forward_reference(model, torch.zeros((1, count)), None, length, chunk_size) + assert not model.calls diff --git a/unit_tests/benchmark/service/test_cache_capacity.py b/unit_tests/benchmark/service/test_cache_capacity.py new file mode 100644 index 0000000000..01bc5cceab --- /dev/null +++ b/unit_tests/benchmark/service/test_cache_capacity.py @@ -0,0 +1,62 @@ +import json + +import pytest + +from test.benchmark.service.cache_capacity import plan_vllm_capacity + + +def metrics(**overrides): + fields = dict( + engine="0", + num_gpu_blocks="12345", + kv_cache_max_concurrency=str(12345 / 258), + kv_cache_size_tokens=str(int(12345 / 258 * 4096)), + block_size="16", + mamba_block_size="256", + mamba_cache_mode="align", + mamba_ssm_cache_dtype="float32", + ) + fields.update(overrides) + labels = ",".join(f"{key}={json.dumps(value)}" for key, value in fields.items()) + return f"# TYPE vllm:cache_config_info gauge\nvllm:cache_config_info{{{labels}}} 1\n" + + +def test_hybrid_capacity_does_not_use_attention_block_size_alone(): + plan = plan_vllm_capacity(metrics(), 4096, 16384) + assert plan["blocks_per_full_request"] == 258 + assert plan["num_gpu_blocks_override"] == 1032 + assert plan["cache_geometry"]["mamba_cache_mode"] == "align" + + +def test_unverified_pool_growth_is_not_planned(): + with pytest.raises(ValueError, match="only shrinks"): + plan_vllm_capacity(metrics(), 4096, 196608) + + +@pytest.mark.parametrize("text", ["", metrics() + metrics(engine="1")]) +def test_missing_or_ambiguous_engines_cannot_yield_capacity_plan(text): + with pytest.raises(ValueError, match="exactly one"): + plan_vllm_capacity(text, 4096, 16384) + + +@pytest.mark.parametrize( + "fields", + [ + dict(num_gpu_blocks="None"), + dict(num_gpu_blocks="0"), + dict(kv_cache_max_concurrency="nan"), + dict(kv_cache_max_concurrency="inf"), + dict(kv_cache_size_tokens="9"), + dict(num_gpu_blocks="12346"), + dict(kv_cache_max_concurrency="0"), + ], +) +def test_uninitialized_or_inconsistent_runtime_metrics_are_rejected(fields): + with pytest.raises(ValueError): + plan_vllm_capacity(metrics(**fields), 4096, 16384) + + +@pytest.mark.parametrize("length,budget", [(True, 16384), (0, 16384), (4096, -1), (4096, 16385)]) +def test_invalid_or_partial_context_budgets_are_rejected(length, budget): + with pytest.raises(ValueError): + plan_vllm_capacity(metrics(), length, budget) diff --git a/unit_tests/benchmark/service/test_completion.py b/unit_tests/benchmark/service/test_completion.py new file mode 100644 index 0000000000..6e568ef264 --- /dev/null +++ b/unit_tests/benchmark/service/test_completion.py @@ -0,0 +1,88 @@ +from copy import deepcopy + +import pytest + +from test.benchmark.service.completion import completion_request, completion_tokens + + +def native(ids=(11, 99), reason="stop"): + return dict( + tokens=[dict(id=token) for token in ids], + prompt_tokens=2, + count_output_tokens=len(ids), + finish_reason=reason, + ) + + +def baseline(ids=(11,), reason="stop"): + return dict( + choices=[dict(index=0, token_ids=list(ids), finish_reason=reason)], + usage=dict(prompt_tokens=2, completion_tokens=len(ids)), + ) + + +@pytest.mark.parametrize("framework", ["lightllm", "vllm", "sglang"]) +def test_request_sends_only_prompt_and_decoding_options(framework): + endpoint, payload = completion_request(framework, "model", [1, 2], 3) + if framework == "lightllm": + assert endpoint == "/generate" and payload["inputs"] == [1, 2] + assert payload["parameters"]["print_eos_token"] + assert payload["parameters"]["ignore_eos"] is False + else: + assert endpoint == "/v1/completions" and payload["prompt"] == [1, 2] + assert payload["ignore_eos"] is False and payload["return_token_ids"] + assert "expected_number" not in payload + + +@pytest.mark.parametrize("framework, factory", [("lightllm", native), ("vllm", baseline), ("sglang", baseline)]) +def test_natural_eos_and_budget_exhaustion_remain_distinct(framework, factory): + result = completion_tokens(factory(), framework, [1, 2], 3, [99]) + assert result["finish_reason"] == "stop" and not result["truncated"] + truncated = completion_tokens(factory((11, 12, 13), "length"), framework, [1, 2], 3, [99]) + assert truncated["token_ids"] == [11, 12, 13] and truncated["truncated"] + + +@pytest.mark.parametrize("framework", ["vllm", "sglang"]) +def test_baseline_may_include_or_omit_terminal_eos(framework): + for ids in [(11,), (11, 99)]: + assert completion_tokens(baseline(ids), framework, [1, 2], 3, [99])["token_ids"] == list(ids) + + +@pytest.mark.parametrize("framework, factory", [("lightllm", native), ("vllm", baseline), ("sglang", baseline)]) +@pytest.mark.parametrize("ids, reason", [((99, 11), "stop"), ((11,), "length"), ((11, 12, 99), "length")]) +def test_invalid_termination_cannot_be_scored(framework, factory, ids, reason): + with pytest.raises(ValueError): + completion_tokens(factory(ids, reason), framework, [1, 2], 3, [99]) + + +def test_native_stop_requires_observed_eos(): + with pytest.raises(ValueError, match="declared EOS"): + completion_tokens(native((11,)), "lightllm", [1, 2], 3, [99]) + + +@pytest.mark.parametrize("ids", [[], [True], [-1], [1.0], [11, 12, 13, 99]]) +def test_invalid_output_ids_and_over_budget_are_rejected(ids): + with pytest.raises(ValueError): + completion_tokens(baseline(ids), "vllm", [1, 2], 3, [99]) + + +def test_usage_and_prompt_identity_are_checked(): + original = baseline() + for field, value in [("prompt_tokens", 3), ("completion_tokens", 2), ("completion_tokens", True)]: + body = deepcopy(original) + body["usage"][field] = value + with pytest.raises(ValueError, match="counts"): + completion_tokens(body, "vllm", [1, 2], 3, [99]) + body = deepcopy(original) + body["choices"][0]["prompt_token_ids"] = [1, 3] + with pytest.raises(ValueError, match="exact prompt"): + completion_tokens(body, "vllm", [1, 2], 3, [99]) + body["choices"] *= 2 + with pytest.raises(ValueError, match="one completion"): + completion_tokens(body, "vllm", [1, 2], 3, [99]) + + +@pytest.mark.parametrize("prompt, limit", [([], 3), ([True], 3), ([-1], 3), ([1], True), ([1], 0)]) +def test_bad_requests_fail_before_transport(prompt, limit): + with pytest.raises(ValueError): + completion_request("lightllm", "model", prompt, limit) diff --git a/unit_tests/benchmark/service/test_completion_runner.py b/unit_tests/benchmark/service/test_completion_runner.py new file mode 100644 index 0000000000..9509ca39ca --- /dev/null +++ b/unit_tests/benchmark/service/test_completion_runner.py @@ -0,0 +1,111 @@ +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import json +import threading + +import pytest +import requests + +from test.benchmark.service.completion_runner import CompletionClient + + +@contextmanager +def service(*, overlap=False, bad_count=False): + barrier, second_done = threading.Barrier(2), threading.Event() + finished = [] + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *args): + pass + + def do_POST(self): + payload = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + assert "expected_number" not in payload + index = payload["inputs"][0] + if overlap: + barrier.wait(timeout=5) + if index == 1: + assert second_done.wait(timeout=5) + body = dict( + tokens=[dict(id=10 + index), dict(id=99)], + prompt_tokens=1, + count_output_tokens=1 if bad_count and index == 2 else 2, + finish_reason="stop", + ) + data = json.dumps(body).encode() + self.send_response(200) + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + finished.append(index) + if index == 2: + second_done.set() + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}", finished + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +def cases(): + return [ + dict(dataset_index=17, prompt_ids=[1], max_output_tokens=3, expected_number="secret"), + dict(dataset_index=3, prompt_ids=[2], max_output_tokens=3, expected_number="secret"), + ] + + +def test_actual_concurrent_http_keeps_order_and_sessions_separate(tmp_path): + sessions = [] + + class Session(requests.Session): + def __init__(self): + super().__init__() + self.owner, self.was_closed = threading.get_ident(), False + sessions.append(self) + + def post(self, *args, **kwargs): + assert self.owner == threading.get_ident() and not self.trust_env + return super().post(*args, **kwargs) + + def close(self): + self.was_closed = True + super().close() + + with service(overlap=True) as (url, finished): + client = CompletionClient( + "lightllm", + "model", + url, + 10, + [99], + tmp_path, + concurrency=2, + session_factory=Session, + ) + try: + results = list(client.results(cases())) + finally: + client.close() + assert finished == [2, 1] + assert [x["token_ids"] for x in results] == [[11, 99], [12, 99]] + assert len(sessions) == 2 and all(x.was_closed for x in sessions) + assert {x.name for x in tmp_path.glob("response-*.txt")} == {"response-17.txt", "response-3.txt"} + + +@pytest.mark.parametrize("concurrency", [1, 2]) +def test_count_failure_is_not_dropped_and_raw_response_survives(tmp_path, concurrency): + with service(bad_count=True) as (url, finished): + client = CompletionClient("lightllm", "model", url, 10, [99], tmp_path, concurrency=concurrency) + try: + results = iter(client.results(cases())) + assert next(results)["token_ids"] == [11, 99] + with pytest.raises(ValueError, match="counts"): + next(results) + finally: + client.close() + assert json.loads((tmp_path / "response-3.txt").read_text())["count_output_tokens"] == 1 diff --git a/unit_tests/benchmark/service/test_stream_cohort.py b/unit_tests/benchmark/service/test_stream_cohort.py new file mode 100644 index 0000000000..48ce7d4a08 --- /dev/null +++ b/unit_tests/benchmark/service/test_stream_cohort.py @@ -0,0 +1,185 @@ +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import json +import threading + +import pytest + +from test.benchmark.service.stream_cohort import StreamingWorkload, distribution, run_cohort, summarize_cohort +from test.benchmark.service.token_stream import request_stream + + +def observation(started, arrivals, ended, input_tokens=2, ids=None): + gaps = [right - left for left, right in zip(arrivals, arrivals[1:])] + return dict( + token_ids=list(range(len(arrivals))) if ids is None else ids, + metrics=dict( + input_tokens=input_tokens, + output_tokens=len(arrivals), + ttft_s=arrivals[0] - started, + tpot_s=sum(gaps) / len(gaps) if gaps else None, + e2e_s=ended - started, + receiver_itl_s=gaps, + ), + timeline=dict(started=started, first_token=arrivals[0], last_token=arrivals[-1], ended=ended), + ) + + +def test_overlapping_requests_use_elapsed_time_and_actual_token_counts(): + results = [observation(1, [1.5, 2.5, 4.5], 5), observation(2, [2.7], 4, input_tokens=4)] + metrics = summarize_cohort(results) + assert metrics["elapsed_s"] == 4 + assert metrics["output_tokens"] == 4 and metrics["output_tokens_per_s"] == 1 + assert metrics["input_tokens_per_s"] == 1.5 + assert metrics["maximum_client_inflight"] == 2 + assert metrics["ttft_s"]["p50"] == pytest.approx(0.6) + assert metrics["e2e_s"]["p50"] == 3 + assert metrics["tpot_s"]["count"] == 1 + assert metrics["receiver_itl_s"]["p50"] == 1.5 + + +def test_cache_aggregates_require_complete_observability(): + results = [observation(1, [1.5], 2), observation(2, [2.5], 3)] + results[0]["metrics"]["cached_input_tokens"] = 1 + partial = summarize_cohort(results) + assert partial["cache_reported_requests"] == 1 + assert partial["cached_input_tokens"] is None and partial["uncached_input_tokens"] is None + results[1]["metrics"]["cached_input_tokens"] = 0 + complete = summarize_cohort(results) + assert complete["cache_reported_requests"] == 2 + assert complete["cached_input_tokens"] == 1 and complete["uncached_input_tokens"] == 3 + + +def test_bursts_and_single_tokens_keep_their_metric_meaning(): + metrics = summarize_cohort([observation(1, [2, 2, 3], 4)]) + assert metrics["receiver_zero_gap_fraction"] == 0.5 + metrics = summarize_cohort([observation(1, [2], 4)]) + assert metrics["tpot_s"] is None and metrics["receiver_itl_s"] is None + assert metrics["receiver_zero_gap_fraction"] is None + + +def test_percentiles_are_interpolated_and_invalid_samples_rejected(): + metrics = distribution([4, 2, 1, 3]) + assert metrics["p50"] == 2.5 and metrics["p95"] == pytest.approx(3.85) + assert distribution([]) is None + for values in ([float("nan")], [float("inf")], [-1]): + with pytest.raises(ValueError): + distribution(values) + + +@pytest.mark.parametrize( + "field,value", + [ + ("started", float("nan")), + ("started", 3), + ("ended", 1), + ("first_token", 4), + ("last_token", 1), + ], +) +def test_invalid_timelines_cannot_produce_throughput(field, value): + result = observation(1, [2, 3], 4) + result["timeline"][field] = value + with pytest.raises(ValueError): + summarize_cohort([result]) + + +@pytest.mark.parametrize( + "field,value", + [("output_tokens", 99), ("output_tokens", True), ("input_tokens", 0), ("receiver_itl_s", []), ("ttft_s", 0)], +) +def test_inconsistent_counts_or_timing_cannot_pass(field, value): + result = observation(1, [2, 3], 4) + result["metrics"][field] = value + with pytest.raises(ValueError): + summarize_cohort([result]) + + +def test_reference_mismatch_invalidates_entire_cohort_and_keeps_evidence(tmp_path): + tasks = [StreamingWorkload("ok", [11, 12], [7, 8]), StreamingWorkload("bad", [11, 12], [7, 8])] + + def request(task, path): + ids = [7, 8] if task.name == "ok" else [7, 9] + path.write_text("raw response") + return observation(1, [2, 3], 4, ids=ids) + + report = run_cohort(tasks, 2, request, tmp_path) + assert report["status"] == "failed" and report["metrics"] is None + assert [record["status"] for record in report["records"]] == ["passed", "failed"] + assert report["records"][1]["result"]["token_ids"] == [7, 9] + assert (tmp_path / "000001.sse").read_text() == "raw response" + + +def test_transport_failure_cannot_be_hidden_by_successful_requests(tmp_path): + task = StreamingWorkload("error", [11], [7]) + + def request(task, path): + path.write_text("partial response") + raise TimeoutError("stream stalled") + + report = run_cohort([task], 1, request, tmp_path) + assert report["status"] == "failed" and report["metrics"] is None + assert report["records"][0]["error_type"] == "TimeoutError" + + +@pytest.mark.parametrize("concurrency", [0, -1, 2, True]) +def test_invalid_concurrency_is_rejected(concurrency, tmp_path): + with pytest.raises(ValueError): + run_cohort([StreamingWorkload("x", [1], [2])], concurrency, lambda *_: None, tmp_path) + + +def test_real_chunked_sse_closed_loop_cohort(tmp_path): + pair = threading.Barrier(2) + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, *_): + pass + + def do_POST(self): + body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + assert body["prompt"] == [11, 12] and body["return_token_ids"] is True + pair.wait(timeout=5) + self.send_response(200) + self.send_header("content-type", "text/event-stream") + self.send_header("transfer-encoding", "chunked") + self.end_headers() + events = [ + {"choices": [{"index": 0, "token_ids": [7, 8]}]}, + {"choices": [{"index": 0, "token_ids": [9], "finish_reason": "length"}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + for payload in [*(json.dumps(event) for event in events), "[DONE]"]: + chunk = ("data: " + payload + "\n\n").encode() + self.wfile.write(f"{len(chunk):x}\r\n".encode() + chunk + b"\r\n") + self.wfile.flush() + self.wfile.write(b"0\r\n\r\n") + self.wfile.flush() + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + task = StreamingWorkload("fixture", [11, 12], [7, 8, 9]) + + def request(task, path): + return request_stream( + f"http://127.0.0.1:{server.server_port}", + "vllm", + "fixture", + list(task.prompt_ids), + len(task.expected_ids), + 5, + path, + ) + + try: + report = run_cohort([task] * 4, 2, request, tmp_path) + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + assert report["status"] == "passed" + assert report["metrics"]["requests"] == 4 + assert report["metrics"]["output_tokens"] == 12 + assert report["metrics"]["maximum_client_inflight"] == 2 + assert len(list(tmp_path.glob("*.sse"))) == 4 diff --git a/unit_tests/benchmark/service/test_stream_timing.py b/unit_tests/benchmark/service/test_stream_timing.py new file mode 100644 index 0000000000..bca4c755e3 --- /dev/null +++ b/unit_tests/benchmark/service/test_stream_timing.py @@ -0,0 +1,79 @@ +import copy +import json + +import pytest + +from test.benchmark.service.stream_timing import streaming_completion_times, summarize_streaming_times + + +def stream(text="hello"): + origin = 10_000_000_000 + return dict( + status=200, + start_perf_ns=origin, + metadata=dict(request_start_ns=1_000_000, request_end_ns=3_001_000_000, was_cancelled=False), + responses=[ + dict( + perf_ns=origin + 500_000_000, + packets=[dict(name="data", value=json.dumps(dict(choices=[dict(text=text)])))], + ), + dict( + perf_ns=origin + 1_000_000_000, + packets=[dict(name="data", value=json.dumps(dict(choices=[dict(text=text)])))], + ), + dict( + perf_ns=origin + 2_000_000_000, + packets=[dict(name="data", value=json.dumps(dict(choices=[dict(text="")])))], + ), + dict(perf_ns=origin + 3_000_000_000, packets=[dict(name="data", value="[DONE]")]), + ], + ) + + +def test_invisible_tokens_do_not_shorten_http_latency(): + assert streaming_completion_times(stream()) == dict( + http_completion_ms=3000, + first_visible_text_ms=500, + last_visible_text_ms=1000, + done_after_last_visible_text_ms=2000, + ) + + +def test_empty_text_retains_http_timing_without_invented_text_latency(): + empty = stream("") + empty["metadata"]["request_end_ns"] += 1_000_000_000 + empty["responses"][-1]["perf_ns"] += 1_000_000_000 + result = summarize_streaming_times([stream(), empty]) + assert result["requests"] == 2 and result["empty_text_responses"] == 1 + assert result["http_completion_all_requests_ms"]["p95"] == 3950 + assert result["http_completion_content_requests_ms"]["p95"] == 3000 + assert result["last_visible_text_ms"]["p95"] == 1000 + assert summarize_streaming_times([empty])["last_visible_text_ms"] is None + + +@pytest.mark.parametrize( + "failure", + ["no_done", "duplicate_done", "after_done", "out_of_order", "clock", "cancelled", "http", "float_timestamp"], +) +def test_rejects_invalid_stream_boundaries(failure): + record = stream() + if failure == "no_done": + record["responses"].pop() + elif failure == "duplicate_done": + record["responses"].append(copy.deepcopy(record["responses"][-1])) + elif failure == "after_done": + response = copy.deepcopy(record["responses"][0]) + response["perf_ns"] = record["responses"][-1]["perf_ns"] + record["responses"].append(response) + elif failure == "out_of_order": + record["responses"][1]["perf_ns"] = record["responses"][0]["perf_ns"] - 1 + elif failure == "clock": + record["metadata"]["request_end_ns"] += record["start_perf_ns"] + elif failure == "cancelled": + record["metadata"]["was_cancelled"] = True + elif failure == "http": + record["status"] = 500 + elif failure == "float_timestamp": + record["responses"][0]["perf_ns"] = float(record["responses"][0]["perf_ns"]) + with pytest.raises(ValueError): + streaming_completion_times(record) diff --git a/unit_tests/benchmark/service/test_token_stream.py b/unit_tests/benchmark/service/test_token_stream.py new file mode 100644 index 0000000000..4ef3e19069 --- /dev/null +++ b/unit_tests/benchmark/service/test_token_stream.py @@ -0,0 +1,287 @@ +import json + +import pytest + +from test.benchmark.service.token_stream import SSEDecoder, TokenStream + + +def test_sse_handles_split_utf8_crlf_and_multiline_json(): + decoder = SSEDecoder() + raw = ': heartbeat\r\ndata: {"text": "中",\r\ndata: "token_ids": [7, 8]}\r\n\r\ndata: [DONE]\n\n'.encode() + events = [] + for byte in raw: + events.extend(decoder.feed(bytes([byte]))) + decoder.finish() + assert json.loads(events[0]) == {"text": "中", "token_ids": [7, 8]} + assert events[1] == "[DONE]" + + +@pytest.mark.parametrize("raw", [b'data: {"x": 1}', b'data: {"x": 1}\n', b'data: "\xe4']) +def test_truncated_sse_cannot_pass(raw): + decoder = SSEDecoder() + decoder.feed(raw) + with pytest.raises((ValueError, UnicodeDecodeError)): + decoder.finish() + + +@pytest.mark.parametrize("framework", ["vllm", "sglang"]) +def test_burst_token_count_and_arrival_metrics(framework): + stream = TokenStream(framework, [11, 12], 3, 1.0) + stream.record(json.dumps({"choices": [{"index": 0, "text": "ab", "token_ids": [7, 8]}]}), 1.5) + stream.record( + json.dumps({"choices": [{"index": 0, "text": "c", "token_ids": [9], "finish_reason": "length"}]}), 1.8 + ) + stream.record(json.dumps({"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}), 1.9) + stream.record("[DONE]", 1.9) + result = stream.summarize(2.0) + assert result["output_tokens"] == 3 and result["events"] == 3 + assert result["ttft_s"] == 0.5 + assert result["tpot_s"] == pytest.approx(0.15) + assert result["receiver_itl_s"] == pytest.approx([0, 0.3]) + + +def test_native_empty_text_is_still_an_output_token(): + stream = TokenStream("lightllm", [11], 1, 1.0) + stream.record( + json.dumps( + { + "token": {"id": 7, "text": "", "prompt_tokens": 1, "count_output_tokens": 1}, + "finished": True, + "finish_reason": "length", + } + ), + 1.2, + ) + assert stream.summarize(1.3)["output_tokens"] == 1 + + +@pytest.mark.parametrize("framework", ["lightllm", "vllm", "sglang"]) +@pytest.mark.parametrize("cached", [None, 0, 1, 2]) +def test_optional_cache_counts_preserve_unknown_and_zero(framework, cached): + stream = TokenStream(framework, [11, 12], 1, 1.0) + if framework == "lightllm": + body = dict(token=dict(id=7, prompt_tokens=2, count_output_tokens=1), finished=True, finish_reason="length") + if cached is not None: + body["token"]["prompt_cache_len"] = cached + else: + body = dict( + choices=[dict(index=0, token_ids=[7], finish_reason="length")], + usage=dict(prompt_tokens=2, completion_tokens=1), + ) + if cached is not None: + body["usage"]["prompt_tokens_details"] = dict(cached_tokens=cached) + stream.record(json.dumps(body), 1.2) + if framework != "lightllm": + stream.record("[DONE]", 1.3) + result = stream.summarize(1.4) + assert result["cached_input_tokens"] == cached + assert result["uncached_input_tokens"] == (None if cached is None else 2 - cached) + + +@pytest.mark.parametrize("cached", [-1, True, 3, 1.0]) +def test_invalid_cached_usage_cannot_pass(cached): + stream = TokenStream("vllm", [11, 12], 1, 1.0) + with pytest.raises(ValueError, match="Cached input"): + stream.record( + json.dumps( + dict( + choices=[], + usage=dict(prompt_tokens=2, completion_tokens=0, prompt_tokens_details=dict(cached_tokens=cached)), + ) + ), + 1.1, + ) + + +def test_cached_usage_cannot_change_within_request(): + stream = TokenStream("vllm", [11, 12], 1, 1.0) + stream.record( + json.dumps( + dict( + choices=[], + usage=dict(prompt_tokens=2, completion_tokens=0, prompt_tokens_details=dict(cached_tokens=1)), + ) + ), + 1.1, + ) + with pytest.raises(ValueError, match="changed"): + stream.record( + json.dumps( + dict( + choices=[], + usage=dict(prompt_tokens=2, completion_tokens=0, prompt_tokens_details=dict(cached_tokens=0)), + ) + ), + 1.2, + ) + + +@pytest.mark.parametrize("cold", [False, True]) +def test_native_transport_requests_observable_details_and_explicit_cold_policy(monkeypatch, tmp_path, cold): + import requests + from test.benchmark.service.token_stream import request_stream + + class Response: + status_code = 200 + headers = {"content-type": "text/event-stream", "transfer-encoding": "chunked"} + + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def iter_content(self, chunk_size): + body = dict( + token=dict(id=7, prompt_tokens=2, count_output_tokens=1, prompt_cache_len=0 if cold else 1), + finished=True, + finish_reason="length", + ) + yield ("data:" + json.dumps(body) + "\n\n").encode() + + class Session: + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def post(self, url, json, **kwargs): + assert json["parameters"]["return_details"] is True + assert json["parameters"].get("disable_prompt_cache", False) is cold + return Response() + + monkeypatch.setattr(requests, "Session", Session) + result = request_stream( + "http://localhost", "lightllm", "model", [11, 12], 1, 10, tmp_path / "response.sse", disable_prompt_cache=cold + ) + assert result["metrics"]["cached_input_tokens"] == (0 if cold else 1) + + +def test_baseline_request_cannot_silently_ignore_native_cache_policy(tmp_path): + from test.benchmark.service.token_stream import request_stream + + with pytest.raises(ValueError, match="only implemented for native"): + request_stream( + "http://localhost", "vllm", "model", [11], 1, 10, tmp_path / "response.sse", disable_prompt_cache=True + ) + + +@pytest.mark.parametrize( + "body", + [ + {"choices": [{"index": 0, "text": "x"}]}, + {"choices": [{"index": 0, "token_ids": [True]}]}, + {"choices": [{"index": 0, "token_ids": [7, 8, 9]}]}, + {"choices": [{"index": 0, "token_ids": [7], "finish_reason": "stop"}]}, + {"choices": [{"index": 0, "prompt_token_ids": [99]}]}, + {"choices": [], "usage": {"prompt_tokens": 99, "completion_tokens": 0}}, + {"error": {"message": "aborted"}}, + {"choices": [{"index": 0, "token_ids": [-1]}]}, + {"choices": [{"index": 0, "token_ids": 0}]}, + ], +) +def test_invalid_stream_events_are_rejected(body): + stream = TokenStream("vllm", [11], 2, 1.0) + with pytest.raises(ValueError): + stream.record(json.dumps(body), 1.1) + + +def test_truncated_generation_and_missing_done_cannot_be_reported(): + stream = TokenStream("vllm", [11], 1, 1.0) + with pytest.raises(ValueError): + stream.record("[DONE]", 1.1) + with pytest.raises(ValueError): + stream.summarize(1.2) + stream.record(json.dumps({"choices": [{"index": 0, "token_ids": [7], "finish_reason": "length"}]}), 1.2) + with pytest.raises(ValueError): + stream.summarize(1.3) + with pytest.raises(ValueError): + stream.record(json.dumps({"choices": [{"index": 0, "token_ids": [8]}]}), 1.3) + + +def test_control_events_also_require_monotonic_timestamps(): + stream = TokenStream("vllm", [11], 1, 1.0) + stream.record(json.dumps({"choices": []}), 1.5) + with pytest.raises(ValueError): + stream.record(json.dumps({"choices": []}), 1.4) + stream.record(json.dumps({"choices": [{"index": 0, "token_ids": [7], "finish_reason": "length"}]}), 1.6) + stream.record("[DONE]", 1.7) + with pytest.raises(ValueError): + stream.summarize(1.65) + + +@pytest.mark.parametrize("framework", ["vllm", "sglang"]) +@pytest.mark.parametrize("cache_salt", [None, "isolated-prefix", "x" * 128]) +def test_transport_uses_chunk_boundaries_and_saves_raw_after_timing(monkeypatch, tmp_path, framework, cache_salt): + import requests + import time + from test.benchmark.service.token_stream import request_stream + + bodies = [ + {"choices": [{"index": 0, "text": "ab", "token_ids": [7, 8]}]}, + {"choices": [{"index": 0, "token_ids": [9], "finish_reason": "length"}]}, + {"choices": [], "usage": {"prompt_tokens": 1, "completion_tokens": 3}}, + ] + chunks = [("data: " + json.dumps(body) + "\n\n").encode() for body in bodies] + chunks.append(b"data: [DONE]\n\n") + + class Response: + status_code = 200 + headers = {"content-type": "text/event-stream", "transfer-encoding": "chunked"} + + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def iter_content(self, chunk_size): + assert chunk_size is None + yield from chunks + + class Session: + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def post(self, url, json, **kwargs): + assert json["prompt"] == [11] and json["return_token_ids"] is True + assert json["stream_options"]["include_usage"] is True + if cache_salt is None: + assert "cache_salt" not in json + else: + assert json["cache_salt"] == cache_salt + return Response() + + monkeypatch.setattr(requests, "Session", Session) + timestamps = iter([1.0, 1.5, 1.8, 1.9, 1.9, 2.0]) + monkeypatch.setattr(time, "perf_counter", lambda: next(timestamps)) + path = tmp_path / "response.sse" + result = request_stream("http://localhost", framework, "model", [11], 3, 10, path, cache_salt=cache_salt) + assert result["token_ids"] == [7, 8, 9] + assert result["metrics"]["ttft_s"] == 0.5 + assert result["metrics"]["e2e_s"] == 1.0 + assert result["timeline"] == {"started": 1.0, "first_token": 1.5, "last_token": 1.8, "ended": 2.0} + assert path.read_bytes() == b"".join(chunks) + + +@pytest.mark.parametrize("cache_salt", ["", "x" * 129, 1, True, ["salt"], "x@y", "x/y", "x\\y", "x\x00y"]) +def test_invalid_cache_salt_is_rejected_before_transport(tmp_path, cache_salt): + from test.benchmark.service.token_stream import request_stream + + with pytest.raises(ValueError, match="Cache salt must"): + request_stream( + "http://localhost", "vllm", "model", [11], 1, 10, tmp_path / "response.sse", cache_salt=cache_salt + ) + + +def test_native_cache_salt_cannot_be_silently_ignored(tmp_path): + from test.benchmark.service.token_stream import request_stream + + with pytest.raises(ValueError, match="only implemented for baseline"): + request_stream( + "http://localhost", "lightllm", "model", [11], 1, 10, tmp_path / "response.sse", cache_salt="salt" + ) diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index 10c52c8cea..86e62cb671 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -253,6 +253,8 @@ def build_hybrid_resources( small_page_capacity: int, draft_full_state: FullAttentionState | None = None, tp_rank: int = 0, + ssm_state_mode: str = "native", + ssm_replay_config: dict | None = None, ) -> AttentionStateManager: """测试按既定几何参数组装通用管理器，不依赖旧模型管理器的构造/绑定流程。""" from lightllm.common.attention_state import AttentionStateManager @@ -282,6 +284,9 @@ def build_hybrid_resources( request_capacity=request_capacity, mtp_step=mtp_step, device="cuda", + ssm_state_mode=ssm_state_mode, + ssm_replay_config=ssm_replay_config, + gate_lower_bound=config.gate_lower_bound if glm else -5.0, ) bounded_state.initialize_page_pools( big_page_capacity=big_page_capacity, 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 15d29bd85d..fb9ce7ba5b 100644 --- a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py +++ b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py @@ -9,6 +9,7 @@ from lightllm.server.core.objs.start_args_type import StartArgs +@pytest.mark.parametrize("mode", ["native", "replay"]) @pytest.mark.parametrize("draft_layers", [0, 1, 3]) @pytest.mark.parametrize( "layer_types,expected_indices,linear_layers", @@ -41,6 +42,7 @@ def test_hybrid_layer_mapping_and_mtp_states( expected_indices: tuple[int, ...], linear_layers: int, draft_layers: int, + mode: str, ) -> None: model_config = dict( num_hidden_layers=len(layer_types), @@ -56,7 +58,13 @@ def test_hybrid_layer_mapping_and_mtp_states( full_attention_interval=4, layer_types=layer_types, ) - args = StartArgs(data_type="bfloat16", linear_att_ssm_data_type="float32") + args = StartArgs( + data_type="bfloat16", + linear_att_ssm_data_type="float32", + mtp_step=2, + ssm_state_mode=mode, + ssm_replay_config={"cache_len": 8} if mode == "replay" else {}, + ) config = Qwen3NextResourceConfig.from_model_config(model_config, args) logical_layer_ids = tuple(range(len(layer_types))) (full_state,) = Qwen3NextTpPartModel.create_full_att_states( @@ -70,6 +78,11 @@ def test_hybrid_layer_mapping_and_mtp_states( model_logical_layer_ids=logical_layer_ids, device="meta", ) + if mode == "native": + stride = bounded_state.mtp_step + 1 + assert bounded_state.ssm_state.shape[1] == bounded_state.request_capacity * stride + if mode == "replay": + assert bounded_state.capacity == 8 manager = AttentionStateManager(request_capacity=1) manager.register_state(model_role="main", state=full_state) manager.register_state(model_role="main", state=bounded_state) diff --git a/unit_tests/common/att_resources/test_page_pool_storage.py b/unit_tests/common/att_resources/test_page_pool_storage.py index 032f7674eb..5eb5ba52cc 100644 --- a/unit_tests/common/att_resources/test_page_pool_storage.py +++ b/unit_tests/common/att_resources/test_page_pool_storage.py @@ -58,9 +58,9 @@ def test_shared_registration_lives_until_last_view( unregistered = [] unregister = utils._unregister_page_pool_storage - def record_unregister(ptr: int, device_id: int) -> None: - unregistered.append(ptr) - unregister(ptr=ptr, device_id=device_id) + def record_unregister(addresses: tuple[int, ...], device_id: int) -> None: + unregistered.extend(addresses) + unregister(addresses=addresses, device_id=device_id) monkeypatch.setattr(utils, "_unregister_page_pool_storage", record_unregister) tensor = utils.create_page_pool_tensor(size=(4096,), dtype=torch.float32) @@ -75,6 +75,92 @@ def record_unregister(ptr: int, device_id: int) -> None: assert unregistered == [ptr] +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_chunked_shared_pages_keep_dma_boundaries_and_lifetime(cpu_args, monkeypatch, dtype) -> None: + cpu_args.enable_dp_prompt_cache_fetch = True + monkeypatch.setattr(utils, "_HOST_REGISTRATION_BYTES", 16384) + runtime = torch.cuda.cudart() + registered, released = [], [] + + class RecordingRuntime: + def cudaHostRegister(self, address, size, flags): + status = runtime.cudaHostRegister(address, size, flags) + if int(status) == 0: + registered.append((address, size)) + return status + + def cudaHostUnregister(self, address): + released.append(address) + return runtime.cudaHostUnregister(address) + + monkeypatch.setattr(torch.cuda, "cudart", lambda: RecordingRuntime()) + tensor = utils.create_page_pool_tensor(size=(7, 3, 512), dtype=dtype) + cell_bytes = tensor[0].numel() * tensor.element_size() + assert len(registered) > 1 + assert all((address - tensor.data_ptr()) % cell_bytes == 0 for address, _ in registered) + assert all(size <= 16384 for _, size in registered) + destination = torch.empty((3, 5, 512), dtype=dtype, device="cuda") + for index in range(tensor.shape[0]): + tensor[index].fill_(index + 1) + destination[:, 1].copy_(tensor[index], non_blocking=True) + torch.cuda.synchronize() + assert torch.equal(destination[:, 1].cpu(), tensor[index]) + + view = tensor[-1] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + destination[:, 1].copy_(view, non_blocking=True) + view.fill_(17) + del tensor + gc.collect() + assert not released + assert view.is_shared() and view.is_pinned() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(destination[:, 1].cpu(), view) + graph.reset() + del view + gc.collect() + assert released == [address for address, _ in reversed(registered)] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_private_storage_keeps_whole_tensor_dma(monkeypatch) -> None: + monkeypatch.setattr(utils, "_HOST_REGISTRATION_BYTES", 16384) + tensor = torch.full((7, 3, 512), 7, dtype=torch.float32) + utils.register_page_pool_tensor(tensor) + destination = tensor.to("cuda", non_blocking=True) + torch.cuda.synchronize() + assert torch.equal(destination.cpu(), tensor) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_shared_pool_partial_registration_failure_rolls_back(cpu_args, monkeypatch) -> None: + cpu_args.enable_dp_prompt_cache_fetch = True + monkeypatch.setattr(utils, "_HOST_REGISTRATION_BYTES", 16384) + runtime = torch.cuda.cudart() + registered, released = [], [] + + class FailingRuntime: + def cudaHostRegister(self, address, size, flags): + if registered: + return 1 + status = runtime.cudaHostRegister(address, size, flags) + if int(status) == 0: + registered.append(address) + return status + + def cudaHostUnregister(self, address): + released.append(address) + return runtime.cudaHostUnregister(address) + + monkeypatch.setattr(torch.cuda, "cudart", lambda: FailingRuntime()) + with pytest.raises(torch.cuda.CudaError): + utils.create_page_pool_tensor(size=(7, 3, 512), dtype=torch.float32) + assert registered and released == registered + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_gpu_page_pool_keeps_device_allocation(cpu_args: EasyDict) -> None: cpu_args.enable_dp_prompt_cache_fetch = True diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 8ec77c6f4d..22b561b822 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -74,6 +74,7 @@ def test_decode_state_initializes_normal_layout(): 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 + assert state.b1_mtp_cu_q_seq_len_int64 is None @pytest.mark.parametrize("role", ["main", "draft0"]) @@ -91,6 +92,7 @@ def test_decode_state_initializes_fixed_mtp_layout(role: str): 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.b1_mtp_cu_q_seq_len_int64, torch.tensor([0, 3, 6], dtype=torch.int64)) 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( @@ -124,6 +126,7 @@ def build_params(**kwargs): state.init_state() assert state.b1_mtp_cu_q_seq_len is expected_cu_q_seq_len + torch.testing.assert_close(state.b1_mtp_cu_q_seq_len_int64, expected_cu_q_seq_len.to(torch.int64)) 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( @@ -140,6 +143,89 @@ def build_params(**kwargs): assert build_calls[0]["hold_req_id"] == 5 +def test_mtp_recurrent_layers_reuse_converted_offsets(monkeypatch): + import lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp as conv_mtp + + state, infer_state = _create_decode_state( + draft_step=2, + dynamic_layout=False, + b_req_idx=torch.tensor([0, 0, 0], dtype=torch.int32), + req_to_mtp_state_index=torch.tensor([0, 0], dtype=torch.int32), + ) + state.init_state() + state.backend.activation = "silu" + state.backend._rearrange_mixed_qkv = lambda mixed, decode: (mixed, mixed, mixed) + offsets = [] + monkeypatch.setattr(conv_mtp, "causal_conv1d_update", lambda mixed, *args, **kwargs: mixed) + + def recurrent(**kwargs): + offsets.append(kwargs["cu_seqlens"]) + return kwargs["q"], None + + monkeypatch.setattr(gdn, "mtp_fused_recurrent_gated_delta_rule", recurrent) + weight = SimpleNamespace( + linear_conv1d=SimpleNamespace(mm_param=SimpleNamespace(weight=None), bias=None), + linear_A_log=SimpleNamespace(weight=None), + linear_dt_bias=SimpleNamespace(weight=None), + ) + tensor = torch.zeros((3, 1)) + for _ in range(2): + state._gdn_mtp_kernel(tensor, tensor, tensor, tensor, tensor, state.infer_state, weight) + assert len(offsets) == 2 + assert all(value is state.b1_mtp_cu_q_seq_len_int64 for value in offsets) + torch.testing.assert_close(offsets[0], torch.tensor([0, 3], dtype=torch.int64)) + + +@pytest.mark.parametrize( + "device", + ["cpu", pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required"))], +) +def test_graph_state_copy_refreshes_cached_mtp_offsets(monkeypatch, device): + layouts = [ + torch.tensor([0, 2, 3, 3], dtype=torch.int32, device=device), + torch.tensor([0, 1, 3, 3], dtype=torch.int32, device=device), + ] + current = [0] + + def build_params(**kwargs): + return ( + layouts[current[0]], + torch.tensor([0, 1, 2], dtype=torch.int32, device=device), + torch.ones(3, dtype=torch.int32, device=device), + ) + + monkeypatch.setattr(gdn, "build_dynamic_mtp_linear_att_state_params", build_params) + + def create(): + state, infer_state = _create_decode_state( + draft_step=2, + dynamic_layout=True, + b_req_idx=torch.tensor([0, 0, 1] if current[0] == 0 else [0, 1, 1], dtype=torch.int32, device=device), + req_to_mtp_state_index=torch.zeros(3, dtype=torch.int32, device=device), + ) + state.init_state() + return state + + captured = create() + cached = captured.b1_mtp_cu_q_seq_len_int64 + graph = None + if device == "cuda": + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + observed = cached.clone() + current[0] = 1 + fresh = create() + captured.copy_for_decode_cuda_graph(fresh) + assert captured.b1_mtp_cu_q_seq_len_int64 is cached + expected = torch.tensor([0, 1, 3, 3], dtype=torch.int64, device=device) + torch.testing.assert_close(cached, expected) + assert captured.b1_mtp_cu_q_seq_len.dtype == torch.int32 + if graph is not None: + graph.replay() + torch.testing.assert_close(observed, expected) + graph.reset() + + def test_linear_backend_is_abstract_and_triton_supplies_chunk_kernel(): assert inspect.isabstract(gdn.LinearAttBackend) backend = object.__new__(TritonLinearAttBackend) @@ -177,6 +263,34 @@ def test_missing_linear_backend_arg_uses_auto_selection(monkeypatch, auto_linear 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_qwen35_linear_prefill_att_backend_class(index=1) + decode_backend = linear_create_utils.get_qwen35_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_flashqla_validation_loads_linear_cache_config(monkeypatch): linear_config = SimpleNamespace( tp_world_size=2, 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 new file mode 100644 index 0000000000..fb5168f661 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py @@ -0,0 +1,243 @@ +import pytest +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update + + +def ReplaySSMExecutor(*args, **kwargs): + kda = kwargs.pop("kda", False) + if kda: + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner + else: + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner + + return _history_owner(*args, **kwargs) + + +def ssm_snapshot(backend, req): + return backend.accepted_ssm(req) + + +from types import SimpleNamespace +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import replay_layer + + +def make(n, width): + hold = n + 3 + ids = list(reversed(range(n))) + [hold] + lengths = [1 + i % width for i in range(n)] + [1] + starts = [] + reqs = [] + mtps = [] + for r, l in zip(ids, lengths): + starts.append(len(reqs)) + reqs.extend([r] * l) + mtps.extend(range(l)) + tensor = lambda x: torch.tensor(x, device="cuda", dtype=torch.int32) + args = [tensor(starts), tensor(reqs), tensor(mtps), tensor([1] * len(reqs)), width] + ssm_backend = ReplayGDNState.__new__(ReplayGDNState) + ssm_backend.cursors = tensor([7] * (hold + 1)) + ssm_backend.hold_req_id = hold + ssm_backend.capacity = max(4, triton.next_power_of_2(width)) + ssm_backend.verify_width = width + manager = SimpleNamespace() + manager.mtp_state_index = tensor([-9] * (hold + 1)) + manager._ssm_executor = ssm_backend + return manager, args, (ids, lengths, starts) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("replay", [False, True]) +def test_request_counts_share_compiled_kernel(monkeypatch, replay): + from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import _commit_prefix, commit_ssm_prefix + + reqs = torch.arange(512, dtype=torch.int32, device="cuda") + lengths = reqs % 4 + 1 + indices = torch.zeros(513, dtype=torch.int32, device="cuda") + cursors = torch.zeros_like(indices) if replay else None + hashes = set() + run = _commit_prefix.run + + def record(*args, **kwargs): + kernel = run(*args, **kwargs) + hashes.add(kernel.hash) + return kernel + + monkeypatch.setattr(_commit_prefix, "run", record) + for count in (1, 2, 7, 16, 23, 24, 25, 31, 32, 33, 255, 256, 257, 511): + indices.fill_(-1) + if replay: + cursors.zero_() + commit_ssm_prefix(indices, reqs[:count], lengths[:count], cursors) + assert torch.equal(indices[:count], lengths[:count] - 1) + assert (indices[count:] == -1).all() + if replay: + assert torch.equal(cursors[:count], lengths[:count]) + assert not cursors[count:].any() + assert len(hashes) == 1 + + +def accept_unfused(m, args): + update(m.mtp_state_index, *args) + m._ssm_executor.accept_updates(args[1][args[0].long()], m.mtp_state_index) + + +def accept_fused(m, args): + reqs = args[1][args[0].long()] + ends = torch.cat((args[0][1:], args[0].new_full((1,), args[3].numel()))) + cumulative = torch.cat((args[3].new_zeros(1), args[3].cumsum(0))) + lengths = cumulative[ends.long()] - cumulative[args[0].long()] + m._ssm_executor.commit_accepted(m.mtp_state_index, reqs, lengths) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_acceptance_fusion_matches_original_and_cpu_under_graph(): + checks = 0 + for width in [1, 2, 4, 8, 16]: + for n in [1, 7, 33, 128]: + m, args, (ids, lengths, starts) = make(n, width) + ref, _, _ = make(n, width) + # Compile both paths before graph capture, then reset before replay. + accept_unfused(ref, args) + accept_fused(m, args) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + accept_fused(m, args) + expected_index = [-9] * len(m.mtp_state_index) + expected_cursor = [7] * len(expected_index) + m.mtp_state_index.fill_(-9) + ref.mtp_state_index.fill_(-9) + m._ssm_executor.cursors.fill_(7) + ref._ssm_executor.cursors.fill_(7) + for step in range(24): + mask = [] + for i, (r, l) in enumerate(zip(ids, lengths)): + count = (step + i) % (l + 1) + mask.extend([1] * count + [0] * (l - count)) + expected_index[r] = count - 1 + if r != m._ssm_executor.hold_req_id: + expected_cursor[r] += count if width > 1 else 1 + args[3].copy_(torch.tensor(mask, device="cuda", dtype=torch.int32)) + accept_unfused(ref, args) + if step % 2: + graph.replay() + else: + accept_fused(m, args) + assert torch.equal(m.mtp_state_index, ref.mtp_state_index) + assert torch.equal(m._ssm_executor.cursors, ref._ssm_executor.cursors) + assert m.mtp_state_index.tolist() == expected_index + assert m._ssm_executor.cursors.tolist() == expected_cursor + checks += 1 + + class _IndexOnly: + def commit_accepted(self, indices, req_indices, accepted_lengths): + commit_ssm_prefix(indices, req_indices, accepted_lengths) + + m._ssm_executor = _IndexOnly() + accept_fused(m, args) + assert m.mtp_state_index.tolist() == expected_index + + print("ACCEPT_EQUIVALENCE_CHECKS=" + str(checks), flush=True) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, compute_mode): + torch.manual_seed(93) + initial = torch.randn(2, 8, 4, 32, 32, device="cuda").to(dtype) * 0.01 + ssm_backends = [ + ReplaySSMExecutor( + initial.clone(), + 8, + 4, + torch.bfloat16, + num_key_heads=2, + compute_mode=compute_mode, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + for _ in range(2) + ] + managers = [] + for ssm_backend in ssm_backends: + manager = SimpleNamespace() + manager._ssm_executor = ssm_backend + manager.mtp_state_index = torch.zeros(8, dtype=torch.int32, device="cuda") + managers.append(manager) + reqs = torch.full((4,), 7, dtype=torch.int32, device="cuda") + cu = torch.tensor([0, 3, 6, 9, 12], dtype=torch.int32, device="cuda") + token_reqs = torch.full((12,), 7, dtype=torch.int32, device="cuda") + mtp_indexes = torch.arange(12, dtype=torch.int32, device="cuda") % 3 + accepted_mask = torch.ones(12, dtype=torch.int32, device="cuda") + q = torch.randn(1, 12, 2, 32, dtype=torch.bfloat16, device="cuda") + k = torch.randn_like(q) + v = torch.randn(1, 12, 4, 32, dtype=torch.bfloat16, device="cuda") + a = torch.full((12, 4), -3.0, dtype=torch.bfloat16, device="cuda") + b = torch.randn_like(a) + alog = torch.zeros(4, device="cuda") + args = [cu[:-1], token_reqs, mtp_indexes, accepted_mask, 4] + + def cycle(i): + ssm_backend = ssm_backends[i] + ssm_backend.prepare_decode(reqs, cu) + outputs = [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, alog, reqs, cu)[0] for layer in range(2)] + (accept_unfused if i == 0 else accept_fused)(managers[i], args) + return outputs + + graphs, outputs = [], [] + for i in range(2): + cycle(i) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = cycle(i) + graphs.append(graph) + outputs.append(output) + for iteration in range(40): + ids = [1, 3, 6] if iteration % 2 else [6, 1, 3] + lengths = [4, 4, 3] if iteration % 3 == 0 else [3, 3, 3] + lengths.append(12 - sum(lengths)) + ids.append(7) + starts, tokens, indexes, masks = [0], [], [], [] + for i, (req, length) in enumerate(zip(ids, lengths)): + count = (iteration + i) % (length + 1) if req != 7 else 0 + tokens.extend([req] * length) + indexes.extend(range(length)) + masks.extend([1] * count + [0] * (length - count)) + starts.append(len(tokens)) + for tensor, values in [ + (reqs, ids), + (cu, starts), + (token_reqs, tokens), + (mtp_indexes, indexes), + (accepted_mask, masks), + ]: + tensor.copy_(tensor.new_tensor(values)) + for tensor in (q, k, v, b): + tensor.normal_() + for graph in graphs: + graph.replay() + for left, right in zip(outputs[0], outputs[1]): + assert torch.equal(left, right) + assert torch.equal(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) + assert torch.equal(managers[0].mtp_state_index, managers[1].mtp_state_index) + if iteration % 7 == 0: + assert torch.equal(ssm_snapshot(ssm_backends[0], 3), ssm_snapshot(ssm_backends[1], 3)) + if iteration % 11 == 10: + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(reqs[:3]) + assert torch.equal(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state) + if iteration == 19: + for manager in managers: + manager._ssm_executor.clear_history(1) + manager._ssm_executor.ssm_state[:, 1].zero_() + manager.mtp_state_index[1] = 0 + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(reqs[:3]) + assert torch.equal(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state) + assert torch.equal(ssm_backends[0].ssm_state[:, [0, 2, 4, 5, 7]], initial[:, [0, 2, 4, 5, 7]]) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d.py new file mode 100644 index 0000000000..1bbc9627b0 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d.py @@ -0,0 +1,404 @@ +import builtins +from collections import defaultdict + +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( + causal_conv1d_fn as dispatch_prefill, + causal_conv1d_update as dispatch_update, +) +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_triton import ( + _prefill_configs, + _update_configs, + causal_conv1d_fn, + causal_conv1d_update, +) +from lightllm.common.triton_utils import autotuner + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _prefill_reference(x, weight, bias, states, lengths, indices, initial, activation): + """Independent grouped-convolution reference, including cached prefixes.""" + width = weight.shape[-1] + output = x.clone() + start = 0 + for length, slot, use_initial in zip(lengths, indices, initial): + if length and slot != -1: + prefix = states[slot] if use_initial else torch.zeros_like(states[slot]) + extended = torch.cat((prefix, x[:, start : start + length]), dim=-1) + with torch.backends.cudnn.flags(allow_tf32=False): + convolved = F.conv1d( + extended.float().unsqueeze(0), + weight.float().unsqueeze(1), + bias.float() if bias is not None else None, + groups=x.shape[0], + ).squeeze(0) + output[:, start : start + length] = F.silu(convolved) if activation else convolved + states[slot] = extended[:, -(width - 1) :] + start += length + return output + + +def _assert_output(actual, expected): + if actual.dtype == torch.float32: + torch.testing.assert_close(actual, expected, rtol=2e-5, atol=2e-6) + else: + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=2e-3) + + +def _check_prefill_ragged_prefix(dtype, width, layout, run_config=None): + torch.manual_seed(42) + dim = 131 # Exercise masked channels and unaligned strides. + lengths = [1, 0, 3, 19, 67] + tokens = sum(lengths) + if layout == "projection": + x = torch.randn(tokens, dim + 32, device="cuda", dtype=dtype)[:, :dim].T + else: + x = torch.randn(dim, tokens * 2, device="cuda", dtype=dtype)[:, ::2] + weight = torch.randn(dim, width * 2, device="cuda", dtype=dtype)[:, ::2] * 0.2 + bias = torch.randn(dim * 2, device="cuda", dtype=dtype)[::2] if layout == "projection" else None + activation = "silu" if bias is not None else None + state_storage = torch.randn(6, dim, width + 2, device="cuda", dtype=dtype) + original_storage = state_storage.clone() + states = state_storage[..., : width - 1] + ref_states = states.clone() + indices = [3, 1, -1, 4, 0] + initial = [True, False, False, True, False] + starts = torch.tensor([0, 1, 1, 4, 23, 90], device="cuda", dtype=torch.int32) + expected = _prefill_reference(x, weight, bias, ref_states, lengths, indices, initial, activation) + actual = causal_conv1d_fn( + x, + weight, + bias, + query_start_loc=starts, + cache_indices=torch.tensor(indices, device="cuda"), + has_initial_state=torch.tensor(initial, device="cuda"), + conv_states=states, + activation=activation, + max_seqlen=max(lengths), + run_config=run_config, + ) + _assert_output(actual, expected) + torch.testing.assert_close(states, ref_states, rtol=0, atol=0) + torch.testing.assert_close(state_storage[..., width - 1 :], original_storage[..., width - 1 :], rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("width", [2, 3, 4]) +@pytest.mark.parametrize("layout", ["projection", "sequence"]) +def test_prefill_ragged_prefix_and_strided_state(dtype, width, layout): + _check_prefill_ragged_prefix(dtype, width, layout) + + +@pytest.mark.parametrize("run_config", _prefill_configs()) +def test_prefill_autotune_configurations(run_config): + _check_prefill_ragged_prefix(torch.bfloat16, 4, "projection", run_config) + + +def test_prefill_new_requests_ignore_stale_state(): + x = torch.randn(2, 65, 33, device="cuda") + weight = torch.randn(65, 4, device="cuda") + states = torch.randn(2, 65, 3, device="cuda") + with torch.backends.cudnn.flags(allow_tf32=False): + expected = F.silu(F.conv1d(F.pad(x, (3, 0)), weight[:, None], groups=65)) + actual = causal_conv1d_fn( + x.transpose(0, 1).reshape(65, -1), + weight, + query_start_loc=torch.tensor([0, 33, 66], device="cuda", dtype=torch.int32), + cache_indices=torch.arange(2, device="cuda", dtype=torch.int32), + has_initial_state=torch.zeros(2, device="cuda", dtype=torch.bool), + conv_states=states, + activation="swish", + ) + _assert_output(actual, expected.transpose(0, 1).reshape(65, -1)) + torch.testing.assert_close(states, x[..., -3:], rtol=0, atol=0) + + +def _update_reference(x, states, weight, bias, indices, cache_lengths, activation): + output = x.clone() + input_3d = x.unsqueeze(-1) if x.ndim == 2 else x + output_3d = output.unsqueeze(-1) if output.ndim == 2 else output + width, state_len = weight.shape[-1], states.shape[-1] + for batch, slot in enumerate(indices): + if slot == -1: + continue + for token in range(input_3d.shape[-1]): + if cache_lengths is None: + history = states[slot, :, -(width - 1) :].clone() + states[slot] = torch.cat((states[slot, :, 1:], input_3d[batch, :, token : token + 1]), dim=-1) + else: + cursor = cache_lengths[batch] + token + history = torch.stack( + [states[slot, :, (cursor - width + 1 + j) % state_len] for j in range(width - 1)], dim=-1 + ) + states[slot, :, cursor % state_len] = input_3d[batch, :, token] + values = torch.cat((history, input_3d[batch, :, token : token + 1]), dim=-1).double() + # Round after each FP32 FMA, independently of Triton's implementation. + accumulator = bias.float().clone() if bias is not None else torch.zeros(input_3d.shape[1], device=x.device) + for tap in range(width): + accumulator = (accumulator.double() + values[:, tap] * weight[:, tap].double()).float() + output_3d[batch, :, token] = F.silu(accumulator) if activation else accumulator + return output + + +def _check_update_buffers(dtype, width, circular, tokens, run_config=None): + torch.manual_seed(17) + dim = 131 + x = torch.randn(3, dim + 16, tokens * 2, device="cuda", dtype=dtype)[:, :dim, ::2] + if tokens == 1: + x = x.squeeze(-1) + weight = torch.randn(dim, width, device="cuda", dtype=dtype) * 0.2 + bias = torch.randn(dim * 2, device="cuda", dtype=dtype)[::2] + state_len = width + 1 + storage = torch.randn(5, dim, state_len + 3, device="cuda", dtype=dtype) + original = storage.clone() + states = storage[..., :state_len] + ref_states = states.clone() + indices = [2, -1, 0] + cache_lengths = [0, 3, 13] if circular else None + activation = "silu" if tokens == 1 else None + expected = _update_reference(x, ref_states, weight, bias, indices, cache_lengths, activation) + actual = causal_conv1d_update( + x, + states, + weight, + bias, + activation, + torch.tensor(cache_lengths, device="cuda", dtype=torch.int32) if circular else None, + torch.tensor(indices, device="cuda", dtype=torch.int32), + run_config=run_config, + ) + _assert_output(actual, expected) + torch.testing.assert_close(states, ref_states, rtol=0, atol=0) + torch.testing.assert_close(storage[..., state_len:], original[..., state_len:], rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("width", [2, 3, 4]) +@pytest.mark.parametrize("circular,tokens", [(False, 1), (False, 7), (True, 1), (True, 7)]) +def test_update_linear_and_ring_buffers(dtype, width, circular, tokens): + _check_update_buffers(dtype, width, circular, tokens) + + +@pytest.mark.parametrize("run_config", _update_configs()) +@pytest.mark.parametrize("circular", [False, True]) +def test_update_autotune_configurations(run_config, circular): + _check_update_buffers(torch.bfloat16, 4, circular, 7, run_config) + + +@pytest.mark.parametrize("mode", ["prefill", "linear", "circular"]) +def test_autotune_and_cached_config_update_state_once(monkeypatch, tmp_path, mode): + torch.manual_seed(13) + tuner = causal_conv1d_fn if mode == "prefill" else causal_conv1d_update + monkeypatch.setattr(autotuner, "get_triton_autotune_level", lambda: autotuner.AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(tuner, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(tuner, "cached_configs", {}) + monkeypatch.setattr(tuner, "fast_match_configs", defaultdict(dict)) + monkeypatch.setattr(tuner, "warmuped_configs_set", set()) + configs = tuner.configs_gen_func() + monkeypatch.setattr(tuner, "configs_gen_func", lambda: [configs[0], configs[-1]]) + + dim, state_len = 65, 3 if mode == "prefill" else 5 + storage = torch.randn(4, dim, state_len + 2, device="cuda", dtype=torch.bfloat16) + original = storage.clone() + states = storage[..., :state_len] + reference_states = states.clone() + weight = torch.randn(dim, 4, device="cuda", dtype=states.dtype) * 0.2 + indices = torch.tensor([3, 1], device="cuda", dtype=torch.int32) + if mode == "prefill": + x = torch.randn(3, dim + 16, device="cuda", dtype=states.dtype)[:, :dim].T + starts = torch.tensor([0, 1, 3], device="cuda", dtype=torch.int32) + initial = torch.tensor([True, True], device="cuda") + + def run(): + return tuner( + x, + weight, + query_start_loc=starts, + cache_indices=indices, + has_initial_state=initial, + conv_states=states, + max_seqlen=2, + ) + + def reference(): + return _prefill_reference(x, weight, None, reference_states, [1, 2], [3, 1], [True, True], "silu") + + else: + x = torch.randn(2, dim, device="cuda", dtype=states.dtype) + lengths = [0, 13] if mode == "circular" else None + cache_lengths = torch.tensor(lengths, device="cuda", dtype=torch.int32) if lengths is not None else None + + def run(): + return tuner(x, states, weight, activation="silu", cache_seqlens=cache_lengths, conv_state_indices=indices) + + def reference(): + return _update_reference(x, reference_states, weight, None, [3, 1], lengths, "silu") + + clone_mutated_args = tuner._mutate_args_clone + scratch_states = [] + + def check_clones(args, kwargs): + new_args, new_kwargs, original_states, cloned_states = clone_mutated_args(args, kwargs) + assert len(cloned_states) == 1 and original_states[0] is states + assert cloned_states[0].data_ptr() != states.data_ptr() + scratch_states.append(cloned_states[0]) + return new_args, new_kwargs, original_states, cloned_states + + monkeypatch.setattr(tuner, "_mutate_args_clone", check_clones) + expected = reference() + with autotuner.Autotuner.autotune_warmup(tuner.kernel_type): + actual = run() + assert len(scratch_states) == 2 # One protected copy per candidate. + assert len(list(tmp_path.glob("*.json"))) == 1 + assert all(config is not None for cached in tuner.cached_configs.values() for config in cached.values()) + _assert_output(actual, expected) + torch.testing.assert_close(states, reference_states, rtol=0, atol=0) + + # Historical-config warmup must also use protected state, including when + # the first cache load is captured in a CUDA graph. + monkeypatch.setattr(autotuner, "get_triton_autotune_level", lambda: autotuner.AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG) + monkeypatch.setattr(tuner, "cached_configs", {}) + monkeypatch.setattr(tuner, "fast_match_configs", defaultdict(dict)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = run() + expected = reference() + graph.replay() + _assert_output(output, expected) + torch.testing.assert_close(states, reference_states, rtol=0, atol=0) + torch.testing.assert_close(storage[..., state_len:], original[..., state_len:], rtol=0, atol=0) + assert len(scratch_states) > 2 + + +def test_prefill_graph_replays_with_changed_lengths_and_prefixes(): + torch.manual_seed(1) + x = torch.randn(25, 96, device="cuda", dtype=torch.bfloat16).T + weight = torch.randn(96, 4, device="cuda", dtype=x.dtype) * 0.2 + original = torch.randn(2, 96, 3, device="cuda", dtype=x.dtype) + states = original.clone() + starts = torch.tensor([0, 1, 25], dtype=torch.int32, device="cuda") + initial = torch.tensor([True, False], device="cuda") + indices = torch.arange(2, device="cuda", dtype=torch.int32) + + def run(): + return causal_conv1d_fn( + x, + weight, + query_start_loc=starts, + cache_indices=indices, + has_initial_state=initial, + conv_states=states, + max_seqlen=25, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = run() + for lengths, flags in [([1, 24], [True, False]), ([25, 0], [False, True]), ([2, 23], [True, True])]: + starts.copy_(torch.tensor([0, lengths[0], 25], dtype=starts.dtype, device="cuda")) + initial.copy_(torch.tensor(flags, device="cuda")) + states.copy_(original) + reference_states = original.clone() + expected = _prefill_reference(x, weight, None, reference_states, lengths, [0, 1], flags, "silu") + graph.replay() + _assert_output(output, expected) + torch.testing.assert_close(states, reference_states, rtol=0, atol=0) + + +def test_dispatch_without_sgl_kernel_or_vllm(monkeypatch): + import_fn = builtins.__import__ + + def without_optional_kernels(name, *args, **kwargs): + if name.split(".")[0] in ("sgl_kernel", "vllm"): + raise ModuleNotFoundError(name) + return import_fn(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", without_optional_kernels) + x = torch.randn(13, 65, dtype=torch.bfloat16, device="cuda").T + weight = torch.randn(65, 4, dtype=x.dtype, device="cuda") * 0.2 + states = torch.zeros(1, 65, 3, dtype=x.dtype, device="cuda") + ref_states = states.clone() + starts = torch.tensor([0, 13], dtype=torch.int32, device="cuda") + expected = _prefill_reference(x, weight, None, ref_states, [13], [0], [False], "silu") + _assert_output( + dispatch_prefill( + x, + weight, + query_start_loc=starts, + conv_states=states, + cache_indices=torch.zeros(1, device="cuda", dtype=torch.int32), + has_initial_state=torch.zeros(1, device="cuda", dtype=torch.bool), + ), + expected, + ) + token = torch.randn(1, 65, dtype=x.dtype, device="cuda") + expected = _update_reference(token, ref_states, weight, None, [0], None, "silu") + _assert_output(dispatch_update(token, states, weight, activation="silu"), expected) + torch.testing.assert_close(states, ref_states, rtol=0, atol=0) + + +@pytest.mark.parametrize("tokens", [1, 3]) +def test_sgl_update(tokens): + pytest.importorskip("sgl_kernel") + torch.manual_seed(29) + x = torch.randn(3, 64, tokens, dtype=torch.bfloat16, device="cuda") + if tokens == 1: + x = x.squeeze(-1) + states = torch.randn(5, 64, 3, dtype=x.dtype, device="cuda") + ref_states = states.clone() + weight = torch.randn(64, 4, dtype=x.dtype, device="cuda") * 0.2 + bias = torch.randn(64, dtype=x.dtype, device="cuda") + indices = [2, -1, 0] + expected = _update_reference(x, ref_states, weight, bias, indices, None, "silu") + actual = dispatch_update( + x, + states, + weight, + bias, + activation="silu", + conv_state_indices=torch.tensor(indices, dtype=torch.int32, device="cuda"), + ) + assert actual.shape == x.shape + _assert_output(actual, expected) + torch.testing.assert_close(states, ref_states, rtol=0, atol=0) + + +def test_bfloat16_products_accumulate_in_float32(): + # BF16 operands whose product cannot be represented in BF16. With a unit + # residual the premature product rounding also changes the BF16 output. + x = torch.tensor([[1.0078125, 1.0]], device="cuda", dtype=torch.bfloat16) + weight = torch.tensor([[1.0078125, -1.0]], device="cuda", dtype=torch.bfloat16) + expected = torch.tensor([[-1.0078125, 0.01568603515625]], device="cuda", dtype=x.dtype) + actual = causal_conv1d_fn( + x, + weight, + activation=None, + query_start_loc=torch.tensor([0, 2], device="cuda", dtype=torch.int32), + cache_indices=torch.zeros(1, device="cuda", dtype=torch.int32), + has_initial_state=torch.zeros(1, device="cuda", dtype=torch.bool), + conv_states=torch.ones(1, 1, 1, device="cuda", dtype=x.dtype), + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_empty_inputs_do_not_touch_state(): + state = torch.randn(2, 17, 3, device="cuda") + original = state.clone() + weight = torch.randn(17, 4, device="cuda") + output = causal_conv1d_fn( + torch.empty(17, 0, device="cuda"), + weight, + conv_states=state, + query_start_loc=torch.zeros(3, device="cuda", dtype=torch.int32), + cache_indices=torch.arange(2, device="cuda", dtype=torch.int32), + has_initial_state=torch.ones(2, device="cuda", dtype=torch.bool), + ) + assert output.numel() == 0 + assert causal_conv1d_update(torch.empty(2, 17, 0, device="cuda"), state, weight).numel() == 0 + torch.testing.assert_close(state, original, rtol=0, atol=0) 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..9021f50195 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 @@ -6,6 +6,31 @@ from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import causal_conv1d_update +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("activation", [None, "silu"]) +def test_low_precision_products_accumulate_in_fp32(dtype, activation): + torch.manual_seed(2907) + x = torch.randn(6, 512, device="cuda", dtype=dtype) * 0.3 + weight = torch.randn(512, 4, device="cuda", dtype=dtype) * 0.2 + prefix = torch.randn(2, 512, 3, device="cuda", dtype=dtype) * 0.3 + full = torch.cat([prefix, x.view(2, 3, 512).transpose(1, 2)], -1).float() + expected = torch.nn.functional.conv1d(full, weight.float()[:, None], groups=512) + if activation == "silu": + expected = torch.nn.functional.silu(expected) + expected = expected.transpose(1, 2).reshape(6, 512).to(dtype) + actual = causal_conv1d_update( + x.clone(), + torch.nn.functional.pad(prefix, (0, 2)), + weight, + mtp_step=2, + conv_state_indices=torch.tensor([0, 1], device="cuda", dtype=torch.int32), + num_accepted_tokens=torch.ones(2, device="cuda", dtype=torch.int32), + query_start_loc=torch.tensor([0, 3, 6], device="cuda", dtype=torch.int32), + activation=activation, + ) + torch.testing.assert_close(actual, expected, atol=1e-7, rtol=1e-6) + + def causal_conv1d_ref( x: torch.Tensor, conv_state: torch.Tensor, diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_packed_decode.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_packed_decode.py new file mode 100644 index 0000000000..aae63dd572 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_packed_decode.py @@ -0,0 +1,62 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( + fused_recurrent_gated_delta_rule, + fused_recurrent_gated_delta_rule_packed_decode, +) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_packed_decode_matches_materialized_decode(): + torch.manual_seed(17) + batch, key_heads, value_heads = 2, 2, 4 + key_dim = value_dim = 128 + device = "cuda" + dtype = torch.bfloat16 + + mixed_qkv = torch.randn( + batch, + 2 * key_heads * key_dim + value_heads * value_dim, + device=device, + dtype=dtype, + ) + a_raw = torch.randn(batch, value_heads, device=device, dtype=dtype) + b_raw = torch.randn_like(a_raw) + a_log = torch.randn(value_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(value_heads, device=device, dtype=torch.float32) + state = torch.randn(batch + 1, value_heads, key_dim, value_dim, device=device, dtype=dtype) + state_reference = state.clone() + state_packed = state.clone() + state_indices = torch.arange(1, batch + 1, device=device, dtype=torch.int32) + + q = mixed_qkv[:, : key_heads * key_dim].view(batch, 1, key_heads, key_dim) + k = mixed_qkv[:, key_heads * key_dim : 2 * key_heads * key_dim].view(batch, 1, key_heads, key_dim) + v = mixed_qkv[:, 2 * key_heads * key_dim :].view(batch, 1, value_heads, value_dim) + + reference, _ = fused_recurrent_gated_delta_rule( + q, + k, + v, + initial_state=state_reference, + ssm_state_indices=state_indices, + use_qk_l2norm_in_kernel=True, + A_log=a_log, + dt_bias=dt_bias, + a_raw=a_raw, + b_raw=b_raw, + ) + packed, _ = fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv, + a_raw, + b_raw, + a_log, + dt_bias, + key_dim ** -0.5, + state_packed, + state_indices, + use_qk_l2norm_in_kernel=True, + ) + + torch.testing.assert_close(packed, reference, atol=0.02, rtol=0.02) + torch.testing.assert_close(state_packed, state_reference, atol=0.02, rtol=0.02) 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..b280f4fabb 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 @@ -163,5 +163,57 @@ def test_variable_length_cu_seqlens(seed): ) +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("steps", [1, 2, 3, 4]) +def test_match_decode_preserves_every_sequential_snapshot(state_dtype, steps): + torch.manual_seed(3107) + batch, heads, value_heads, width = 2, 4, 12, 128 + count = batch * steps + q = torch.randn(batch, steps, heads, width, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(batch, steps, value_heads, width, device="cuda", dtype=torch.bfloat16) + a = torch.randn(batch, steps, value_heads, device="cuda", dtype=torch.bfloat16) + b = torch.randn_like(a) + indices = torch.arange(count, device="cuda", dtype=torch.int32).view(batch, steps) + accepted = torch.tensor([1, steps], device="cuda", dtype=torch.int32) + initial = torch.randn(count, value_heads, width, width, device="cuda", dtype=state_dtype) * 0.1 + state = initial[indices.gather(1, accepted.long()[:, None] - 1).flatten().long()].clone() + expected_states = initial.clone() + gates = dict(A_log=torch.zeros(value_heads, device="cuda"), dt_bias=torch.zeros(value_heads, device="cuda")) + outputs = [] + for step in range(steps): + output, _ = fused_recurrent_gated_delta_rule( + q=q[:, step : step + 1], + k=k[:, step : step + 1], + v=v[:, step : step + 1], + initial_state=state, + inplace_final_state=True, + ssm_state_indices=torch.arange(batch, device="cuda", dtype=torch.int32), + use_qk_l2norm_in_kernel=True, + a_raw=a[:, step], + b_raw=b[:, step], + **gates, + ) + outputs.append(output) + expected_states.index_copy_(0, indices[:, step].long(), state) + actual, actual_states = mtp_fused_recurrent_gated_delta_rule( + q=q.reshape(1, count, heads, width), + k=k.reshape(1, count, heads, width), + v=v.reshape(1, count, value_heads, width), + initial_state=initial.clone(), + cu_seqlens=torch.arange(batch + 1, device="cuda", dtype=torch.int64) * steps, + ssm_state_indices=indices, + ssm_state_write_indices=indices, + num_accepted_tokens=accepted, + a_raw=a.reshape(count, value_heads), + b_raw=b.reshape(count, value_heads), + run_config={"BV": 32, "num_warps": 4, "num_stages": 1}, + **gates, + ) + expected = torch.cat(outputs, dim=1).reshape_as(actual) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + torch.testing.assert_close(actual_states, expected_states, atol=0, rtol=0) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) 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 new file mode 100644 index 0000000000..8252c445a1 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py @@ -0,0 +1,65 @@ +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, +) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("request_capacity,token_bucket", [(4, 3), (4, 16), (128, 512)]) +def test_dynamic_sequences_stay_within_request_capacity_under_graph(request_capacity, token_bucket): + hold = request_capacity + reqs = torch.full((token_bucket,), hold, dtype=torch.int32, device="cuda") + mtp_index = torch.zeros_like(reqs) + config = ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=1, + key_dim=2, + value_dim=2, + conv_width=3, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + ) + manager = GDNState(config, request_capacity + 1, mtp_step=3) + manager.mtp_state_index.copy_(torch.arange(request_capacity + 1, dtype=torch.int32, device="cuda") % 4) + sequences = min(request_capacity, token_bucket) + + def build(): + return build_dynamic_mtp_linear_att_state_params(reqs, mtp_index, manager.mtp_state_index, hold) + + build() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + cu, ids, counts = build() + assert cu.numel() == token_bucket + 1 + assert ids.numel() == counts.numel() == token_bucket + + # Reuse one graph for full request pools, mixed verify lengths, and HOLD-only padding. + for lengths in ([1] * sequences, [min(4, token_bucket // sequences)] * sequences, [2, 1], []): + lengths = lengths[:request_capacity] + if sum(lengths) > token_bucket: + continue + request_ids = list(reversed(range(len(lengths)))) + host_reqs = [req for req, length in zip(request_ids, lengths) for _ in range(length)] + host_mtp = [j for length in lengths for j in range(length)] + tokens = len(host_reqs) + host_reqs += [hold] * (token_bucket - tokens) + host_mtp += [0] * (token_bucket - tokens) + reqs.copy_(torch.tensor(host_reqs, dtype=torch.int32, device="cuda")) + mtp_index.copy_(torch.tensor(host_mtp, dtype=torch.int32, device="cuda")) + graph.replay() + + expected_cu = [0] + for length in lengths: + expected_cu.append(expected_cu[-1] + length) + expected_cu += [tokens] * (token_bucket + 1 - len(expected_cu)) + expected_ids = request_ids + [hold] * (token_bucket - len(request_ids)) + assert cu.tolist() == expected_cu + assert ids.tolist() == expected_ids + expected_counts = [req % 4 + 1 for req in request_ids] + [1] * (token_bucket - len(request_ids)) + assert counts.tolist() == expected_counts 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 new file mode 100644 index 0000000000..22235598ed --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -0,0 +1,776 @@ +import pytest +import torch + + +def ReplaySSMExecutor(*args, **kwargs): + kda = kwargs.pop("kda", False) + if kda: + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner + else: + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner + + return _history_owner(*args, **kwargs) + + +def ssm_snapshot(backend, req): + return backend.accepted_ssm(req) + + +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs, replay_layer + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +@pytest.mark.parametrize("flush", [False, True], ids=["compress", "export"]) +@pytest.mark.parametrize("capture", [False, True], ids=["eager", "graph"]) +def test_multi_layer_fold_keeps_cursor_until_all_tiles_finish(kda, flush, capture): + state = torch.ones(36, 9, 6, 128, 128, device="cuda") + backend = ReplaySSMExecutor( + state, + capacity=16, + verify_width=4, + kda=kda, + num_key_heads=2, + run_config={"BV": 32, "num_warps": 4, "num_stages": 1}, + ) + backend.raw_keys.zero_() + backend.raw_values.zero_() + backend.gates.fill_(-0.1) + backend.betas.zero_() + backend.cursors[:8] = 13 + reqs = torch.arange(8, device="cuda", dtype=torch.int32) + expected = torch.full_like(state[:, :8], torch.exp(torch.tensor(-1.3)).item()) + targets = [torch.empty_like(state[:, 0]) for _ in range(8)] if flush else None + + def fold(): + if flush: + backend.export_accepted_state(list(range(8)), targets, flush=True) + else: + backend.maybe_compress_prefix(reqs) + + if capture: + fold() + state.fill_(1) + backend.cursors[:8] = 13 + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fold() + state.fill_(1) + backend.cursors[:8] = 13 + graph.replay() + else: + fold() + if flush: + torch.testing.assert_close(torch.stack(targets, dim=1), expected, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(state[:, :8], expected, atol=1e-6, rtol=1e-6) + assert not backend.cursors.any() + + +def reference(q, k, v, a, b, a_log, bias, state, kda=False, gate_lower_bound=-5.0): + q, k, v = q.float(), k.float(), v.float() + q = q / (q.square().sum(-1, keepdim=True) + 1e-6).sqrt() * (q.shape[-1] ** -0.5) + k = k / (k.square().sum(-1, keepdim=True) + 1e-6).sqrt() + q = q.repeat_interleave(v.shape[-2] // q.shape[-2], -2) + k = k.repeat_interleave(v.shape[-2] // k.shape[-2], -2) + if kda: + g = gate_lower_bound * torch.sigmoid(a_log.exp()[:, None] * (a.float().view_as(k) + bias.view_as(k))) + state = state * g.exp()[..., None] + else: + g = -a_log.exp() * torch.nn.functional.softplus(a.float() + bias) + state = state * g.exp()[..., None, None] + d = b.float().sigmoid()[..., None] * (v - torch.einsum("...kv,...k->...v", state, k)) + state = state + k[..., None] * d[..., None, :] + return torch.einsum("...kv,...k->...v", state, q), state + + +@pytest.mark.parametrize("capacity", [4, 8]) +@pytest.mark.parametrize("bv", [8, 32], ids=["small_bv_fallback", "precomputed"]) +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +@pytest.mark.parametrize("activation_dtype", [torch.float16, torch.bfloat16]) +def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state_dtype, kda, activation_dtype): + state = torch.zeros(1, 2, 1, 128, 128, dtype=state_dtype, device="cuda") + ssm_backend = ReplaySSMExecutor( + state, + capacity, + 2, + activation_dtype, + num_key_heads=1, + kda=kda, + compute_mode="precompute", + run_config={"BV": bv, "num_warps": 1}, + ) + reqs = torch.tensor([0, 1], dtype=torch.int32, device="cuda") + cu = torch.tensor([0, 2, 4], dtype=torch.int32, device="cuda") + accepted = torch.tensor([1, 0], dtype=torch.int32, device="cuda") + q = torch.zeros(1, 4, 1, 128, dtype=activation_dtype, device="cuda") + q[0, :2, 0, 0] = 1 + k = torch.zeros_like(q) + k[0, 0, 0, 0] = -0.25 + k[0, 1, 0, 0] = 1 + v = torch.zeros_like(q) + v[0, :2, 0, 0] = 60000 + a = torch.full((4, 128 if kda else 1), -20, dtype=activation_dtype, device="cuda") + b = torch.full((4, 1), 20, dtype=activation_dtype, device="cuda") + alog = torch.zeros(1, device="cuda") + bias = torch.zeros(128 if kda else 1, device="cuda") + expected_state = state[0, 0].float().clone() + for round_idx in range(2): + ssm_backend.prepare_decode(reqs, cu) + # The second round must consume the first round's accepted history. + # For capacity 4 the next span exactly fills it, so no fold should hide + # an overflowed history correction by reconstructing from raw inputs. + assert ssm_backend.cursors[0].item() == 2 * round_idx + assert not state.any() + actual = replay_layer(ssm_backend, 0, q, k, v, a, b, alog, bias, reqs, cu)[0] + expected = [] + for token in range(2): + output, expected_state = reference( + q[0, token], k[0, token], v[0, token], a[token], b[token], alog, bias, expected_state, kda=kda + ) + expected.append(output) + assert torch.isfinite(expected_state).all() + assert torch.isfinite(actual).all() + torch.testing.assert_close( + actual[0, :2].float(), torch.stack(expected).to(actual.dtype).float(), atol=4, rtol=1e-3 + ) + ssm_backend.accept_updates(reqs, accepted) + # The correction is larger than FP16 can represent even though the + # recurrent state, inputs, and attention outputs are all finite. + assert ssm_backend.deltas[0, 0, 0, 1, 0].item() > torch.finfo(torch.float16).max + assert torch.isfinite(ssm_backend.deltas[0, 0, :, :4]).all() + assert ssm_backend.cursors[1].item() == 0 + torch.testing.assert_close(ssm_snapshot(ssm_backend, 0)[0].float(), expected_state, atol=128, rtol=0.01) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_flush_export_updates_checkpoint_and_clears_history(dtype): + state = torch.randn(1, 2, 1, 32, 64, device="cuda", dtype=dtype) * 0.01 + backend = ReplaySSMExecutor( + state, + capacity=4, + verify_width=1, + activation_dtype=torch.bfloat16, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + backend.raw_keys.normal_(0, 0.1) + backend.raw_values.normal_(0, 0.1) + backend.gates.fill_(-0.1) + backend.betas.fill_(0.5) + backend.cursors[0] = 2 + expected = ssm_snapshot(backend, 0).clone() + target = torch.empty_like(expected) + + backend.export_accepted_state([0], [target], flush=True) + torch.cuda.synchronize() + + assert backend.cursors[0].item() == 0 + torch.testing.assert_close(target, expected, atol=0, rtol=0) + torch.testing.assert_close(state[:, 0], expected, atol=0, rtol=0) + + empty_history_target = torch.empty_like(expected) + backend.export_accepted_state([0], [empty_history_target], flush=True) + torch.cuda.synchronize() + torch.testing.assert_close(empty_history_target, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("capacity", [16, 32, 64]) +@pytest.mark.parametrize("width", [1, 3, 5, 16]) +@pytest.mark.parametrize("dims", [(32, 64), (128, 128)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_replay_acceptance_flush_and_graph( + width, + dims, + capacity, + dtype, + run_config=None, + compute_mode="inline", + activation_dtype=torch.bfloat16, + gate_shift=-3, + head_geometry=(2, 4), + kda=False, + gate_lower_bound=-5.0, +): + torch.manual_seed(21) + kdim, vdim = dims + h, hv = head_geometry + slots = 5 + state = (torch.randn(2, slots, hv, kdim, vdim, device="cuda") * 0.01).to(dtype) + ssm_backend = ReplaySSMExecutor( + state, + capacity, + width, + num_key_heads=h, + kda=kda, + gate_lower_bound=gate_lower_bound, + compute_mode=compute_mode, + run_config=run_config, + activation_dtype=activation_dtype, + ) + expected = state.float().clone() + pending = [0, 0] + state_atol = 2e-5 if dtype == torch.float32 else 0.002 + # Independent FP32 reductions can straddle a BF16 rounding boundary. + # Allow one relative BF16 spacing as well as accumulated near-zero error. + state_rtol = 2e-4 if dtype == torch.float32 else torch.finfo(dtype).eps + reqs = torch.tensor([2, 0, slots - 1], device="cuda", dtype=torch.int32) + cu = torch.tensor([0, width, 2 * width - 1, 2 * width - 1], device="cuda", dtype=torch.int32) + if width == 1: + cu = torch.tensor([0, 1, 2, 2], device="cuda", dtype=torch.int32) + tokens = 3 * width + q = torch.randn(1, tokens, h, kdim, device="cuda", dtype=activation_dtype) + k = torch.randn_like(q) + v = torch.randn(1, tokens, hv, vdim, device="cuda", dtype=activation_dtype) + a = torch.randn(tokens, hv * kdim if kda else hv, device="cuda", dtype=activation_dtype) + gate_shift + b = torch.randn(tokens, hv, device="cuda", dtype=activation_dtype) + alog = torch.randn(hv, device="cuda") * 0.1 + bias = torch.randn(hv * kdim if kda else hv, device="cuda") * 0.1 + accepted = torch.zeros(slots, dtype=torch.int32, device="cuda") + + def step(): + ssm_backend.prepare_decode(reqs, cu) + out = [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, bias, reqs, cu)[0] for layer in range(2)] + ssm_backend.accept_updates(reqs, accepted) + return out + + # Compile on disposable state, then capture with the HOLD slot only. + saved = state.clone() + step() + state.copy_(saved) + ssm_backend.cursors.zero_() + reqs.fill_(slots - 1) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = step() + reqs.copy_(torch.tensor([2, 0, slots - 1], device="cuda", dtype=torch.int32)) + for iteration in range(3 * capacity + 5): + counts = [iteration % width + 1, min((iteration + 1) % width + 1, max(width - 1, 1))] + accepted[2] = counts[0] - 1 + accepted[0] = counts[1] - 1 + for seq, req in enumerate([2, 0]): + query_len = width if seq == 0 else max(width - 1, 1) + if pending[seq] + query_len > capacity: + ssm_backend.merge_accepted_updates(torch.tensor([req], dtype=torch.int32, device="cuda")) + expected[:, req] = expected[:, req].to(dtype).float() + pending[seq] = 0 + pending[seq] += counts[seq] + graph.replay() + for layer in range(2): + for seq, req in enumerate([2, 0]): + start = seq * width + length = width if seq == 0 else max(width - 1, 1) + cur = expected[layer, req].clone() + for j in range(length): + t = start + j + out, cur = reference( + q[0, t], + k[0, t], + v[0, t], + a[t], + b[t], + alog, + bias, + cur, + kda=kda, + gate_lower_bound=gate_lower_bound, + ) + torch.testing.assert_close(outputs[layer][0, t].float(), out, atol=0.006, rtol=0.02) + if j + 1 == counts[seq]: + expected[layer, req].copy_(cur) + if iteration == 2 * capacity + 1: + ssm_backend.merge_accepted_updates(reqs[:2]) + expected = expected.to(dtype).float() + pending = [0, 0] + torch.testing.assert_close(state[:, :4].float(), expected[:, :4], atol=state_atol, rtol=state_rtol) + ssm_backend.merge_accepted_updates(reqs[:2]) + torch.testing.assert_close(state.float(), expected.to(dtype).float(), atol=state_atol, rtol=state_rtol) + + +@pytest.mark.parametrize( + "state_dtype,activation_dtype", + [(torch.float32, torch.float32), (torch.float32, torch.bfloat16), (torch.bfloat16, torch.bfloat16)], +) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +@pytest.mark.parametrize("dims", [(32, 64), (128, 128)]) +@pytest.mark.parametrize("width", [1, 4]) +@pytest.mark.parametrize("gate_shift", [-3, -10]) +def test_fold_matches_reference_across_precisions( + state_dtype, + activation_dtype, + compute_mode, + dims, + width, + gate_shift, +): + test_replay_acceptance_flush_and_graph( + width, + dims, + 8, + state_dtype, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + compute_mode=compute_mode, + activation_dtype=activation_dtype, + gate_shift=gate_shift, + ) + + +@pytest.mark.parametrize("capacity,width", [(4, 1), (4, 3), (4, 4), (8, 1), (8, 3), (8, 4), (8, 5)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_short_replay_history_matches_independent_reference(capacity, width, dtype): + test_replay_acceptance_flush_and_graph(width, (128, 128), capacity, dtype) + + +@pytest.mark.parametrize("capacity,width", [(4, 1), (4, 3), (4, 4), (8, 1), (8, 4), (8, 5), (16, 4), (16, 16), (32, 4)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_precomputed_projections_match_independent_reference(capacity, width, dtype): + test_replay_acceptance_flush_and_graph( + width, + (128, 128), + capacity, + dtype, + run_config={"BV": 32, "num_warps": 1}, + compute_mode="precompute", + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_production_gdn_geometry_matches_independent_reference(dtype): + test_replay_acceptance_flush_and_graph( + 4, + (128, 128), + 8, + dtype, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + compute_mode="precompute", + head_geometry=(4, 12), + ) + + +@pytest.mark.parametrize("warps", [1, 4]) +@pytest.mark.parametrize("decay", [-3, -10]) +@pytest.mark.parametrize( + "state_dtype,dtype", + [(torch.float32, torch.float32), (torch.bfloat16, torch.bfloat16), (torch.float32, torch.bfloat16)], +) +@pytest.mark.parametrize("bv", [8, 32]) +@pytest.mark.parametrize("width", [1, 4]) +def test_precomputed_projections_preserve_precision(warps, decay, state_dtype, dtype, width, bv): + torch.manual_seed(35) + state = (torch.randn((1, 3, 4, 128, 128), device="cuda") * 0.1).to(state_dtype) + config = {"BV": bv, "num_warps": warps} + ssm_backends = [ + ReplaySSMExecutor( + state.clone(), + 16 if width == 1 else 8, + width, + dtype, + num_key_heads=2, + compute_mode=mode, + run_config=config, + ) + for mode in ("inline", "precompute") + ] + reqs = torch.tensor([1, 0], dtype=torch.int32, device="cuda") + tokens = 2 if width == 1 else 7 + cu = torch.tensor([0, 1, 2] if width == 1 else [0, 4, 7], dtype=torch.int32, device="cuda") + q = torch.randn((1, tokens, 2, 128), device="cuda", dtype=dtype) + k = torch.randn_like(q) + v = torch.randn((1, tokens, 4, 128), device="cuda", dtype=dtype) + a = torch.randn((tokens, 4), device="cuda", dtype=dtype) + decay + b = torch.randn_like(a) + alog = torch.zeros(4, device="cuda") + accepted = torch.zeros(3, dtype=torch.int32, device="cuda") + for step in range(64): + for tensor in (q, k, v, b): + tensor.normal_() + accepted[0] = step % max(1, width - 1) + accepted[1] = step % width + outputs = [] + step_lengths = (cu[1:] - cu[:-1]).tolist() + if any( + ssm_backends[0].cursors[req] + length > ssm_backends[0].capacity + for req, length in zip(reqs.tolist(), step_lengths) + ): + overflow = [ + req + for req, length in zip(reqs.tolist(), step_lengths) + if ssm_backends[0].cursors[req] + length > ssm_backends[0].capacity + ] + overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(overflow_reqs) + for ssm_backend in ssm_backends: + ssm_backend.prepare_decode(reqs, cu) + outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, alog, reqs, cu)[0]) + ssm_backend.accept_updates(reqs, accepted) + if dtype == torch.float32: + torch.testing.assert_close(outputs[0], outputs[1], atol=2e-6, rtol=2e-5) + else: + torch.testing.assert_close(outputs[0], outputs[1], atol=5e-4, rtol=0.01) + error = torch.linalg.vector_norm(outputs[0].float() - outputs[1].float()) + scale = torch.linalg.vector_norm(outputs[0].float()) + assert error < scale * (torch.finfo(dtype).eps / 2) + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(reqs) + torch.testing.assert_close(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state, atol=0, rtol=0) + + +def test_bf16_commit_rounds_before_the_next_token(): + # 前缀写满后先压回 SSM，下一 token 从长度 0 重写。两次 commit 必须得到同一块 BF16 状态。 + state = torch.full((1, 2, 1, 32, 64), 0.123, device="cuda", dtype=torch.bfloat16) + ssm_backends = [ReplaySSMExecutor(state.clone(), 16, 1, torch.float32) for _ in range(2)] + reqs = torch.tensor([0], device="cuda", dtype=torch.int32) + for ssm_backend in ssm_backends: + ssm_backend.raw_keys.zero_() + ssm_backend.raw_values.zero_() + ssm_backend.gates.fill_(-0.01) + ssm_backend.betas.zero_() + ssm_backend.cursors[0] = 16 + ssm_backend.merge_accepted_updates(reqs) + q = torch.ones((1, 1, 32), device="cuda") + k = torch.zeros_like(q) + v = torch.zeros((1, 1, 64), device="cuda") + a = torch.full((1, 1), -3.0, device="cuda") + b = torch.zeros_like(a) + alog = torch.zeros(1, device="cuda") + bias = torch.zeros_like(alog) + outputs = [] + for ssm_backend in ssm_backends: + assert ssm_backend.cursors[0].item() == 0 + outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, bias, reqs)[0]) + torch.testing.assert_close(outputs[0], outputs[1], atol=1e-7, rtol=1e-6) + torch.testing.assert_close(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state, atol=0, rtol=0) + + +def test_large_bf16_state_precompute_matches_inline(): + # 输出投影用 tf32x3。重建缓存的 k/d 在 bf16 激活下是 fp16，输出允许一个 fp16 单位。 + # 折回 SSM 仍从原始 k/v 在 fp32 中计算，两条路径的 checkpoint 必须一致。 + heads = 16 + config = {"BV": 32, "num_warps": 1, "num_stages": 1} + reqs = torch.tensor([0], dtype=torch.int32, device="cuda") + alog = torch.zeros(heads, device="cuda") + bias = torch.zeros_like(alog) + + def check(capacity, width, steps): + torch.manual_seed(7) + state = (torch.randn((1, 2, heads, 128, 128), device="cuda") * 10).bfloat16() + backends = [ + ReplaySSMExecutor( + state.clone(), + capacity, + width, + torch.bfloat16, + num_key_heads=heads, + compute_mode=mode, + run_config=config, + ) + for mode in ("inline", "precompute") + ] + cu = None if width == 1 else torch.tensor([0, width], dtype=torch.int32, device="cuda") + accepted = torch.tensor([width - 1], dtype=torch.int32, device="cuda") + for step in range(steps): + torch.manual_seed(100 + step) + q = torch.randn((1, width, heads, 128), device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn((1, width, heads, 128), device="cuda", dtype=torch.bfloat16) + a = torch.randn((width, heads), device="cuda", dtype=torch.bfloat16) - 1 + b = torch.randn_like(a) + outputs = [replay_layer(backend, 0, q, k, v, a, b, alog, bias, reqs, cu)[0] for backend in backends] + # 重建用的 k/d 在 bf16 激活下是 fp16。两条路径在存入前的舍入可以差一个 fp16 单位。 + torch.testing.assert_close(outputs[0], outputs[1], atol=5e-4, rtol=5e-4) + for backend in backends: + backend.accept_updates(reqs, accepted, width=width) + backend.maybe_compress_prefix(reqs) + for backend in backends: + backend.merge_accepted_updates(reqs) + torch.testing.assert_close(backends[0].ssm_state, backends[1].ssm_state, atol=0, rtol=0) + + check(4, 1, 16) + check(16, 2, 4) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +@pytest.mark.parametrize("capture_metadata", [False, True]) +def test_dynamic_verify_lengths_fill_history_before_folding(dtype, compute_mode, capture_metadata): + torch.manual_seed(31) + state = (torch.randn(2, 4, 4, 128, 128, device="cuda") * 0.01).to(dtype) + ssm_backend = ReplaySSMExecutor( + state, 8, 4, dtype, num_key_heads=2, compute_mode=compute_mode, run_config={"BV": 32, "num_warps": 1} + ) + expected = state.float().clone() + pending = [0, 0] + reqs = torch.full((3,), 3, dtype=torch.int32, device="cuda") + cu = torch.tensor([0, 4, 8, 8], dtype=torch.int32, device="cuda") + accepted = torch.zeros(4, dtype=torch.int32, device="cuda") + q = torch.randn(1, 8, 2, 128, dtype=dtype, device="cuda") + k = torch.randn_like(q) + v = torch.randn(1, 8, 4, 128, dtype=dtype, device="cuda") + a = torch.full((8, 4), -3.0, dtype=dtype, device="cuda") + b = torch.randn_like(a) + alog = torch.zeros(4, device="cuda") + + def forward(): + return [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, alog, reqs, cu)[0] for layer in range(2)] + + def step(): + positions = ssm_backend.prepare_decode(reqs, cu) + out = forward() + ssm_backend.accept_updates(reqs, accepted) + return positions, out + + step() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + positions = torch.zeros_like(reqs) + with torch.cuda.graph(graph): + if capture_metadata: + positions, outputs = step() + else: + outputs = forward() + lengths = [(3, 4), (2, 1), (3, 2), (1, 4), (4, 3), (0, 2), (2, 0), (4, 4)] + for iteration in range(24): + ids = [0, 1] if iteration < 3 or iteration % 2 == 0 else [1, 0] + widths = lengths[iteration % len(lengths)] + counts = ( + list(widths) if iteration < 3 else [min(width, 1 + (iteration + i) % 4) for i, width in enumerate(widths)] + ) + reqs.copy_(reqs.new_tensor(ids + [3])) + cu.copy_(cu.new_tensor([0, widths[0], sum(widths), 8])) + accepted.fill_(-1) + for req, count in zip(ids, counts): + accepted[req] = count - 1 + for tensor in (q, k, v, b): + tensor.normal_() + before = state.clone() + for req, width in zip(ids, widths): + if pending[req] + width > 8: + ssm_backend.merge_accepted_updates(torch.tensor([req], dtype=torch.int32, device="cuda")) + expected[:, req] = expected[:, req].to(dtype).float() + pending[req] = 0 + old_cursors = [pending[req] for req in ids] + [0] + if not capture_metadata: + positions.copy_(ssm_backend.prepare_decode(reqs, cu)) + graph.replay() + if not capture_metadata: + ssm_backend.accept_updates(reqs, accepted) + assert positions.tolist() == old_cursors + for seq, (req, width, count) in enumerate(zip(ids, widths, counts)): + for layer in range(2): + current = expected[layer, req].clone() + for j in range(width): + t = (0 if seq == 0 else widths[0]) + j + out, current = reference(q[0, t], k[0, t], v[0, t], a[t], b[t], alog, alog, current) + torch.testing.assert_close(outputs[layer][0, t].float(), out, atol=0.006, rtol=0.02) + if j + 1 == count: + expected[layer, req].copy_(current) + pending[req] += count + assert ssm_backend.cursors.tolist() == pending + [0, 0] + for output in outputs: + assert torch.count_nonzero(output[:, sum(widths) :]).item() == 0 + torch.testing.assert_close(state[:, 2:], before[:, 2:], atol=0, rtol=0) + if iteration == 2: + assert pending[0] == 8 + torch.testing.assert_close(state[:, 0], before[:, 0], atol=0, rtol=0) + ssm_backend.merge_accepted_updates(reqs[:2]) + torch.testing.assert_close( + state.float(), + expected.to(dtype).float(), + atol=2e-5 if dtype == torch.float32 else 0.002, + rtol=2e-4 if dtype == torch.float32 else torch.finfo(dtype).eps, + ) + + +def test_snapshot_without_history_copies_active_state(): + state = torch.randn((2, 3, 2, 32, 64), device="cuda", dtype=torch.bfloat16) + ssm_backend = ReplaySSMExecutor(state, 4, 3) + snapshot = ssm_snapshot(ssm_backend, 1) + torch.testing.assert_close(snapshot, state[:, 1], rtol=0, atol=0) + assert ssm_backend.cursors[1].item() == 0 + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("run_config", get_configs(), ids=lambda c: f"bv{c['BV']}-w{c['num_warps']}-s{c['num_stages']}") +def test_autotune_candidates_match_recurrence(dtype, run_config): + test_replay_acceptance_flush_and_graph(3, (32, 64), 8, dtype, run_config) + + +@pytest.mark.parametrize("save_big", [False, True]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_checkpoint_restores_pending_history_and_accepted_conv(runtime_start_args, save_big, dtype): + from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig + from unit_tests.common.test_linear_state import make_state + + torch.manual_seed(42) + state_rtol = 2e-5 if dtype == torch.float32 else torch.finfo(dtype).eps + config = ConvSSMStateConfig( + logical_layer_ids=(0, 1), + global_k_heads=2, + global_v_heads=2, + key_dim=32, + value_dim=64, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=dtype, + ) + state, _, layout = make_state(mode="replay", config=config) + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) + state.mtp_state_index[1] = 1 + state.conv_state.normal_() + state.ssm_state.zero_() + state.raw_keys.normal_() + state.raw_values.normal_() + state.gates.uniform_(-0.1, -0.01) + state.betas.uniform_(0.1, 0.9) + active_before = state.ssm_state[:, 1].clone() + conv_expected = state.conv_state[:, 1, :, 1:4].clone() + block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + kind = "big" if save_big else "small" + pool = state.big_pages_pool if save_big else state.small_pages_pool + for count in (2, 1): + state.cursors[1] = count + expected = active_before.float().clone() + for i in range(count): + key = state.raw_keys[:, 1, :, i].float() + key /= (key.square().sum(-1, keepdim=True) + 1e-6).sqrt() + expected *= state.gates[:, 1, :, i].exp()[..., None, None] + delta = state.betas[:, 1, :, i, None] * ( + state.raw_values[:, 1, :, i].float() - torch.einsum("lhkv,lhk->lhv", expected, key) + ) + expected += key[..., None] * delta[..., None, :] + materialized = ssm_snapshot(state, 1) + torch.testing.assert_close(materialized, expected.to(dtype), rtol=state_rtol, atol=2e-5) + state.save_bounded_att_runtime_to_pages([1], [8], kind, [0]) + state.write_to_pd_page(1, block, 0, 8, layout) + torch.cuda.synchronize() + torch.testing.assert_close(pool.ssm_state[0], materialized.cpu(), rtol=0, atol=0) + torch.testing.assert_close(pool.conv_state[0], conv_expected.cpu(), rtol=0, atol=0) + torch.testing.assert_close(state.ssm_state[:, 1], active_before, rtol=0, atol=0) + assert state.cursors[1].item() == count + for restore in ( + lambda: state.restore_bounded_att_runtime_from_pages([2], [8], kind, [0]), + lambda: state.read_from_pd_page(block, 2, 0, 8, layout), + ): + state.cursors[2] = 3 + state.mtp_state_index[2] = 2 + restore() + assert state.cursors[2].item() == state.mtp_state_index[2].item() == 0 + torch.testing.assert_close(state.ssm_state[:, 2], materialized, rtol=0, atol=0) + torch.testing.assert_close(state.conv_state[:, 2, :, :3], conv_expected, rtol=0, atol=0) + state.reset_bounded_att_state([2]) + assert not state.ssm_state[:, 2].any() + assert not state.conv_state[:, 2].any() + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("capacity,width", [(4, 1), (4, 3), (8, 4), (16, 5), (32, 4), (64, 4)]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +@pytest.mark.parametrize("gate_lower_bound,gate_shift", [(-5.0, 0), (-2.5, -6)]) +@pytest.mark.parametrize("activation_dtype", [torch.float32, torch.bfloat16]) +def test_kda_replay_reference_fold_and_graph( + dtype, capacity, width, compute_mode, gate_lower_bound, gate_shift, activation_dtype +): + test_replay_acceptance_flush_and_graph( + width, + (128, 128), + capacity, + dtype, + run_config={"BV": 32, "num_warps": 2, "num_stages": 1}, + compute_mode=compute_mode, + gate_shift=gate_shift, + activation_dtype=activation_dtype, + head_geometry=(2, 2), + kda=True, + gate_lower_bound=gate_lower_bound, + ) + + +@pytest.mark.parametrize("config", get_configs()) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +def test_kda_tuning_candidates_match_reference(config, compute_mode): + test_replay_acceptance_flush_and_graph( + 3, + (128, 128), + 4, + torch.bfloat16, + run_config=config, + compute_mode=compute_mode, + head_geometry=(2, 2), + kda=True, + gate_shift=-3, + ) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize( + "bucket,live,folding,width,separate", + [ + (12, 12, 12, 4, True), + (13, 13, 13, 4, False), + (255, 255, 127, 4, True), + (256, 256, 128, 4, False), + (257, 257, 128, 4, False), + (256, 256, 128, 3, True), + (256, 12, 12, 4, True), + (256, 13, 4, 4, True), + ], +) +def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, separate, kda): + 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, run_config={"BV": 32, "num_warps": 1}) + reqs = torch.cat( + ( + torch.arange(live, device="cuda", dtype=torch.int32), + torch.full((bucket - live,), live, device="cuda", dtype=torch.int32), + ) + ) + indices = torch.zeros(live + 1, device="cuda", dtype=torch.int32) + expected = initial[0, :live].clone() + alog = torch.zeros(1, device="cuda") + bias = torch.zeros(32 if kda else 1, device="cuda") + for step, query_width in enumerate((4, 4, width)): + if step == 2: + backend.clear_history(reqs[folding:live]) + expected[folding:] = initial[0, folding:live] + cu = (torch.arange(bucket + 1, device="cuda", dtype=torch.int32) * query_width).clamp_max(live * query_width) + q = torch.randn(live, query_width, 1, 32, device="cuda") * 0.2 + k, v = torch.randn_like(q), torch.randn_like(q) + a = torch.full((live, query_width, 32 if kda else 1), -2.0, device="cuda") + b = torch.randn(live, query_width, 1, device="cuda") + if (backend.cursors[:live] + query_width > backend.capacity).any(): + backend.merge_accepted_updates(reqs[:live]) + backend.prepare_decode(reqs, cu) + actual = replay_layer( + backend, + 0, + q.reshape(1, -1, 1, 32), + k.reshape(1, -1, 1, 32), + v.reshape(1, -1, 1, 32), + a.flatten(0, 1), + b.flatten(0, 1), + alog, + bias, + reqs, + cu, + )[0].reshape(live, query_width, 1, 32) + for token in range(query_width): + output, expected = reference( + q[:, token], + k[:, token], + v[:, token], + a[:, token], + b[:, token], + alog, + bias.expand(live, 1, 32) if kda else bias, + expected, + kda=kda, + ) + torch.testing.assert_close(actual[:, token], output, atol=3e-4, rtol=3e-4) + lengths = torch.full((live,), query_width, device="cuda", dtype=torch.int32) + backend.commit_accepted(indices, reqs[:live], lengths) + backend.merge_accepted_updates(reqs[:live]) + torch.testing.assert_close(backend.ssm_state[0, :live], expected, atol=3e-4, rtol=3e-4) + assert torch.equal(backend.ssm_state[0, live], initial[0, live]) 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 new file mode 100644 index 0000000000..7261cf1adb --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py @@ -0,0 +1,473 @@ +import pytest +import torch + + +def ReplaySSMExecutor(*args, **kwargs): + kda = kwargs.pop("kda", False) + if kda: + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner + else: + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner + + return _history_owner(*args, **kwargs) + + +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import replay_layer +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) +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 import replayssm + + kernels = ("_decode_replay", "_advance_cursor") + hashes = {name: set() for name in kernels} + launches = [] + for name in kernels: + kernel = getattr(replayssm, name) + + def record(*args, run=kernel.run, name=name, **kwargs): + compiled = run(*args, **kwargs) + hashes[name].add(compiled.hash) + launches.append(name) + return compiled + + monkeypatch.setattr(kernel, "run", record) + + torch.manual_seed(137) + initial = torch.randn(1, 513, 1, 32, 32, device="cuda") * 0.01 + eager, captured = [ + ReplaySSMExecutor(initial.clone(), 8, 1, torch.float32, kda=kda, run_config={"BV": 32, "num_warps": 1}) + for _ in range(2) + ] + reqs = torch.arange(512, device="cuda", dtype=torch.int32) + q = torch.randn(512, 1, 32, device="cuda") * 0.2 + k, v = torch.randn_like(q), torch.randn_like(q) + a = torch.full((512, 32 if kda else 1), -3.0, device="cuda") + beta = torch.randn(512, 1, device="cuda") + alog, bias = torch.zeros(1, device="cuda"), torch.zeros(a.shape[-1], device="cuda") + + def cycle(backend, ids): + count = ids.numel() + backend.prepare_decode(ids) + out = replay_layer(backend, 0, q[:count], k[:count], v[:count], a[:count], beta[:count], alog, bias, ids)[0] + backend.accept_single_token(ids) + return out + + graphs = {} + for bucket in (32, 512): + ids = reqs[:bucket].clone() + cycle(captured, ids) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = cycle(captured, ids) + graphs[bucket] = graph, ids, out + captured.ssm_state.copy_(initial) + captured.cursors.zero_() + counts = (1, 2, 7, 16, 23, 24, 25, 31, 32, 33, 255, 256, 257, 511) + for count in counts: + graph, ids, output = graphs[32 if count <= 32 else 512] + ids.fill_(captured.hold_req_id) + ids[:count].copy_(reqs[:count]) + for _ in range(8): + live = reqs[:count] + if (eager.cursors[live.long()] + 1 > eager.capacity).any(): + eager.merge_accepted_updates(live) + captured.merge_accepted_updates(ids[:count]) + expected = cycle(eager, reqs[:count]) + calls = len(launches) + graph.replay() + assert len(launches) == calls + torch.testing.assert_close(output[:count], expected, atol=3e-5, rtol=3e-5) + assert torch.equal(captured.cursors, eager.cursors) + torch.testing.assert_close(captured.ssm_state, eager.ssm_state, atol=3e-5, rtol=3e-5) + assert len(hashes["_decode_replay"]) == 1 + assert len(hashes["_advance_cursor"]) == 1 + assert captured.cursors[-1] == 0 + assert torch.equal(captured.ssm_state[:, -1], initial[:, -1]) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("capture_metadata", [False, True]) +@pytest.mark.parametrize("varlen", [False, True]) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_request_manager_accept_with_replay_graph(dtype, capture_metadata, varlen, kda): + torch.manual_seed(827) + config = {"BV": 32, "num_warps": 1, "num_stages": 1} + initial = torch.randn(3, 5, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 + ssm_backends = [ + cls( + initial.clone(), + 8, + 4, + torch.bfloat16, + num_key_heads=2, + kda=kda, + compute_mode="precompute", + run_config=config, + ) + for cls in [ReplaySSMExecutor, ReplaySSMExecutor] + ] + manager = SimpleNamespace() + manager._ssm_executor = ssm_backends[1] + manager.mtp_state_index = torch.zeros(5, device="cuda", dtype=torch.int32) + rows = torch.full((12,), 4, device="cuda", dtype=torch.int32) + mtp = torch.zeros_like(rows) + flags = torch.zeros_like(rows) + starts = torch.zeros(3, device="cuda", dtype=torch.int32) + accepted = torch.full((5,), -1, device="cuda", dtype=torch.int32) + q = torch.randn(1, 12, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, 12, 4, 128, device="cuda", dtype=torch.bfloat16) + a = torch.full((12, 4 * 128 if kda else 4), -3.0, device="cuda", dtype=torch.bfloat16) + beta = torch.randn(a.shape[0], 4, device="cuda", dtype=a.dtype) + alog = torch.zeros(4, device="cuda") + bias = torch.zeros(4 * 128 if kda else 4, device="cuda") + + def inputs(iteration): + ids = [0, 1, 3] if iteration % 2 == 0 else [3, 0, 1] + lengths = [1 + (iteration + j) % 4 for j in range(3)] if varlen else [4] * 3 + counts = [1 + (iteration // 2 + j) % length for j, length in enumerate(lengths)] + offsets, req_rows, indexes, mask = [], [], [], [] + table = [-1] * 5 + for req, length, count in zip(ids, lengths, counts): + offsets.append(len(req_rows)) + req_rows.extend([req] * length) + indexes.extend(range(length)) + mask.extend([int(j < count) for j in range(length)]) + table[req] = count - 1 + logical = len(req_rows) + req_rows.extend([4] * (12 - logical)) + indexes.extend([0] * (12 - logical)) + mask.extend([0] * (12 - logical)) + rows.copy_(rows.new_tensor(req_rows)) + mtp.copy_(mtp.new_tensor(indexes)) + flags.copy_(flags.new_tensor(mask)) + starts.copy_(starts.new_tensor(offsets)) + accepted.copy_(accepted.new_tensor(table)) + return ids, lengths, counts, logical + + def metadata(ssm_backend): + cu, reqs, _counts = build_dynamic_mtp_linear_att_state_params(rows, mtp, manager.mtp_state_index, 4) + return cu, reqs, ssm_backend.prepare_decode(reqs, cu) + + def forward_accept(index, cu, reqs): + ssm_backend = ssm_backends[index] + out = [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu)[0] for layer in range(3)] + if index == 0: + ssm_backend.accept_updates(reqs, accepted) + else: + manager._ssm_executor.commit_accepted( + manager.mtp_state_index, + rows[starts.long()], + accepted[rows[starts.long()]] + 1, + ) + return out + + inputs(0) + graphs, captured, outputs = [], [], [] + for index, ssm_backend in enumerate(ssm_backends): + cu, reqs, pos = metadata(ssm_backend) + forward_accept(index, cu, reqs) + ssm_backend.ssm_state.copy_(initial) + ssm_backend.cursors.zero_() + manager.mtp_state_index.zero_() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + if capture_metadata: + cu, reqs, pos = metadata(ssm_backend) + out = forward_accept(index, cu, reqs) + graphs.append(g) + captured.append((cu, reqs, pos)) + outputs.append(out) + for ssm_backend in ssm_backends: + ssm_backend.ssm_state.copy_(initial) + ssm_backend.cursors.zero_() + manager.mtp_state_index.zero_() + folds = 0 + for iteration in range(32): + ids, lengths, counts, logical = inputs(iteration) + for tensor in [q, k, v, beta]: + tensor.normal_() + overflow = [ + req for req, length in zip(ids, lengths) if ssm_backends[0].cursors[req] + length > ssm_backends[0].capacity + ] + if overflow: + overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(overflow_reqs) + folds += len(overflow) + for index, ssm_backend in enumerate(ssm_backends): + if not capture_metadata: + for dst, src in zip(captured[index], metadata(ssm_backend)): + dst.copy_(src) + graphs[index].replay() + torch.cuda.synchronize() + for expected, actual in zip(captured[0], captured[1]): + assert torch.equal(actual, expected) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) + for expected, actual in zip(outputs[0], outputs[1]): + torch.testing.assert_close( + actual[:, :logical].float(), expected[:, :logical].float(), atol=0.006, rtol=0.02 + ) + torch.testing.assert_close( + ssm_backends[1].ssm_state.float(), + ssm_backends[0].ssm_state.float(), + atol=0.002, + rtol=torch.finfo(dtype).eps, + ) + table = manager.mtp_state_index.cpu().tolist() + for req, count in zip(ids, counts): + assert table[req] == count - 1 + assert table[2] == table[4] == 0 + assert ssm_backends[1].cursors[4].item() == 0 + assert folds > 0 + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("capture_metadata", [False, True]) +@pytest.mark.parametrize("stagger", [False, True]) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_decode_accept_order_preserves_forward_snapshot(dtype, capture_metadata, stagger, kda): + torch.manual_seed(731) + initial = torch.randn(3, 3, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 + ssm_backends = [ + ReplaySSMExecutor( + initial.clone(), + 8, + 1, + dtype, + num_key_heads=2, + kda=kda, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + for _ in range(2) + ] + for ssm_backend in ssm_backends: + ssm_backend.fold_programs = 3 + reqs = torch.tensor([0, 1, 2], device="cuda", dtype=torch.int32) + q = torch.randn(3, 2, 128, device="cuda", dtype=dtype) + k = torch.randn_like(q) + v = torch.randn(3, 4, 128, device="cuda", dtype=dtype) + a = torch.full((3, 4 * 128 if kda else 4), -3.0, device="cuda", dtype=dtype) + beta = torch.randn(a.shape[0], 4, device="cuda", dtype=a.dtype) + alog = torch.zeros(4, device="cuda") + bias = torch.zeros(4 * 128 if kda else 4, device="cuda") + + def forward(ssm_backend): + return [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs)[0] for layer in range(3)] + + def warm(ssm_backend): + ssm_backend.prepare_decode(reqs) + out = forward(ssm_backend) + ssm_backend.accept_updates(reqs) + return out + + graphs, positions, outputs = [], [], [] + for ssm_backend in ssm_backends: + warm(ssm_backend) + ssm_backend.ssm_state.copy_(initial) + ssm_backend.cursors.zero_() + if stagger: + reqs.copy_(reqs.new_tensor([0, 2, 2])) + for _ in range(3): + warm(ssm_backend) + reqs.copy_(reqs.new_tensor([0, 1, 2])) + pos = torch.zeros_like(reqs) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + if capture_metadata: + pos = ssm_backend.prepare_decode(reqs) + out = forward(ssm_backend) + if capture_metadata: + ssm_backend.accept_updates(reqs) + graphs.append(graph) + positions.append(pos) + outputs.append(out) + for iteration in range(24): + reqs.copy_(reqs.new_tensor([0, 1, 2] if iteration % 2 == 0 else [1, 0, 2])) + for tensor in [q, k, v, beta]: + tensor.normal_() + overflow = [int(req) for req in reqs.tolist() if ssm_backends[0].cursors[req] + 1 > ssm_backends[0].capacity] + if overflow: + overflow_reqs = torch.tensor(sorted(set(overflow)), dtype=torch.int32, device="cuda") + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(overflow_reqs) + old = ssm_backends[0].cursors[reqs.long()].clone() + for ssm_backend, graph, pos in zip(ssm_backends, graphs, positions): + if not capture_metadata: + pos.copy_(ssm_backend.prepare_decode(reqs)) + graph.replay() + if not capture_metadata: + ssm_backend.accept_updates(reqs) + torch.cuda.synchronize() + assert torch.equal(positions[0], old) + assert torch.equal(positions[0], positions[1]) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) + for expected, actual in zip(outputs[0], outputs[1]): + torch.testing.assert_close(actual.float(), expected.float(), atol=0.006, rtol=0.02) + torch.testing.assert_close( + ssm_backends[1].ssm_state.float(), + ssm_backends[0].ssm_state.float(), + atol=0.002, + rtol=torch.finfo(dtype).eps, + ) + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_cross_layer_fold_with_state_offset_above_int32(kda): + if torch.cuda.mem_get_info()[0] < 6 * 1024 ** 3: + pytest.skip("large-stride regression requires 6 GiB of free device memory") + torch.manual_seed(791) + shape = (48, 3, 12, 128, 128) + strides = (50_528_256, 196608, 16384, 128, 1) + assert strides[0] < 2 ** 31 < 47 * strides[0] + dense = torch.randn(shape, device="cuda", dtype=torch.bfloat16) * 0.01 + sparse = torch.empty_strided(shape, strides, device="cuda", dtype=dense.dtype) + sparse.copy_(dense) + config = {"BV": 32, "num_warps": 1, "num_stages": 1} + key_heads = 12 if kda else 4 + ssm_backends = [ + ReplaySSMExecutor(state, 8, 4, torch.bfloat16, num_key_heads=key_heads, kda=kda, run_config=config) + for state in [dense, sparse] + ] + q = torch.randn(1, 8, key_heads, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, 8, 12, 128, device="cuda", dtype=torch.bfloat16) + a = torch.full((8, 12 * 128 if kda else 12), -3.0, device="cuda", dtype=torch.bfloat16) + beta = torch.randn(8, 12, device="cuda", dtype=torch.bfloat16) + alog = torch.zeros(12, device="cuda") + bias = torch.zeros(12 * 128 if kda else 12, device="cuda") + reqs = torch.tensor([0, 2], device="cuda", dtype=torch.int32) + cu = reqs.new_tensor([0, 4, 8]) + + def step(ssm_backend, accepted): + overflow = [int(req) for req in reqs.tolist() if ssm_backend.cursors[req] + 4 > ssm_backend.capacity] + if overflow: + ssm_backend.merge_accepted_updates(reqs.new_tensor(overflow)) + pos = ssm_backend.prepare_decode(reqs, cu) + out = [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu)[0] for layer in range(48)] + ssm_backend.accept_updates(reqs, accepted) + return pos, out + + for count in [2, 3]: + accepted = reqs.new_tensor([count - 1, -1, -1]) + for ssm_backend in ssm_backends: + step(ssm_backend, accepted) + reqs.copy_(reqs.new_tensor([0, 1])) + expected_pos, expected_out = step(ssm_backends[0], reqs.new_tensor([0, 0, -1])) + actual_pos, actual_out = step(ssm_backends[1], reqs.new_tensor([0, 0, -1])) + torch.cuda.synchronize() + assert torch.equal(actual_pos, expected_pos) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) + for layer in range(48): + torch.testing.assert_close(actual_out[layer].float(), expected_out[layer].float(), atol=0.006, rtol=0.02) + torch.testing.assert_close( + sparse[layer].float(), dense[layer].float(), atol=0.002, rtol=torch.finfo(torch.bfloat16).eps + ) + + +@pytest.mark.parametrize("batch", [1, 8, 12]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_small_batch_matches_larger_graph_padding(batch, dtype, kda): + torch.manual_seed(1087 + batch) + initial = torch.randn(2, 14, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 + config = {"BV": 32, "num_warps": 1, "num_stages": 1} + ssm_backends = [ + ReplaySSMExecutor(initial.clone(), 8, 4, num_key_heads=2, kda=kda, compute_mode="precompute", run_config=config) + for _ in range(2) + ] + reqs = torch.arange(batch, device="cuda", dtype=torch.int32) + padded = torch.cat((reqs, reqs.new_full((13 - batch,), 13))) + cu = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * 4 + padded_cu = torch.cat((cu, cu.new_full((13 - batch,), batch * 4))) + q = torch.randn(1, batch * 4, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, batch * 4, 4, 128, device="cuda", dtype=torch.bfloat16) + a = torch.full((batch * 4, 4 * 128 if kda else 4), -3.0, device="cuda", dtype=torch.bfloat16) + beta = torch.randn(a.shape[0], 4, device="cuda", dtype=a.dtype) + alog = torch.zeros(4, device="cuda") + bias = torch.zeros(4 * 128 if kda else 4, device="cuda") + accepted = torch.zeros(14, device="cuda", dtype=torch.int32) + graphs, outputs = [], [] + for ssm_backend, ids, offsets in zip(ssm_backends, [reqs, padded], [cu, padded_cu]): + + def step(): + ssm_backend.prepare_decode(ids, offsets) + out = [ + replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, ids, offsets)[0] for layer in range(2) + ] + ssm_backend.accept_updates(ids, accepted) + return out + + step() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = step() + graphs.append(graph) + outputs.append(out) + ssm_backend.ssm_state.copy_(initial) + ssm_backend.cursors.zero_() + for iteration in range(32): + accepted[:batch] = (torch.arange(batch, device="cuda") + iteration) % 4 + for value in [q, k, v, beta]: + value.normal_() + overflow = [idx for idx in range(batch) if ssm_backends[0].cursors[idx] + 4 > ssm_backends[0].capacity] + if overflow: + overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(overflow_reqs) + for graph in graphs: + graph.replay() + for small, padded_out in zip(outputs[0], outputs[1]): + torch.testing.assert_close(small.float(), padded_out.float(), atol=0.006, rtol=0.02) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) + torch.testing.assert_close( + ssm_backends[0].ssm_state.float(), + ssm_backends[1].ssm_state.float(), + atol=0.002, + rtol=torch.finfo(dtype).eps, + ) + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_full_prefix_does_not_address_past_capacity(kda): + torch.manual_seed(4) + capacity = 4 + backend = ReplaySSMExecutor( + torch.randn(1, 3, 2, 32, 32, device="cuda") * 0.01, + capacity, + 1, + torch.bfloat16, + kda=kda, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + reqs = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + q = torch.randn(2, 1, 32, device="cuda") * 0.2 + k, v = torch.randn_like(q), torch.randn_like(q) + a = torch.full((2, 32 if kda else 1), -3.0, device="cuda") + beta = torch.randn(2, 1, device="cuda") + alog, bias = torch.zeros(1, device="cuda"), torch.zeros(a.shape[-1], device="cuda") + backend.ssm_state.zero_() + backend.cursors.zero_() + for name in ("keys", "deltas", "gates", "raw_keys", "raw_values", "betas"): + getattr(backend, name).zero_() + backend.cursors[reqs.long()] = capacity + 2 + replay_layer(backend, 0, q, k, v, a, beta, alog, bias, reqs) + backend.merge_accepted_updates(reqs) + torch.cuda.synchronize() + assert torch.equal(backend.cursors[reqs.long()], torch.zeros(2, device="cuda", dtype=torch.int32)) + assert torch.isfinite(backend.ssm_state).all() + for _ in range(12): + replay_layer(backend, 0, q, k, v, a, beta, alog, bias, reqs) + backend.accept_single_token(reqs) + backend.maybe_compress_prefix(reqs) + torch.cuda.synchronize() + assert int(backend.cursors[reqs.long()].max()) <= capacity + assert torch.isfinite(backend.ssm_state).all() 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 new file mode 100644 index 0000000000..726eec31c4 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py @@ -0,0 +1,82 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning + + +def ReplaySSMExecutor(*args, **kwargs): + kda = kwargs.pop("kda", False) + if kda: + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner + else: + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner + + return _history_owner(*args, **kwargs) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("width", [1, 5]) +def test_rebuild_uses_one_layer_scratch(width): + state = torch.zeros((4, 3, 2, 32, 32), dtype=torch.bfloat16, device="cuda") + ssm_backend = ReplaySSMExecutor(state, 8, width, num_key_heads=2) + tokens = 3 * width + q = torch.ones((1, tokens, 2, 32) if width > 1 else (tokens, 2, 32), dtype=torch.bfloat16, device="cuda") + gates = torch.zeros((tokens, 2), dtype=torch.bfloat16, device="cuda") + a_log = torch.zeros(2, device="cuda") + cu = torch.arange(4, device="cuda", dtype=torch.int32) * width if width > 1 else None + _, kwargs = tuning.rebuild_inputs(ssm_backend, "replay", q, q, q, gates, gates, a_log, a_log, cu_seqlens=cu) + scratch = kwargs["ssm_backend"] + assert scratch is not ssm_backend + assert scratch.ssm_state.shape[0] == 1 + assert scratch.capacity == ssm_backend.capacity + assert torch.count_nonzero(ssm_backend.ssm_state) == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@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 + + state = torch.zeros((2, 4, 2, 32, 32), dtype=torch.bfloat16, device="cuda") + ssm_backend = ReplaySSMExecutor(state, 8, width, num_key_heads=2) + for tensor in vars(ssm_backend).values(): + if isinstance(tensor, torch.Tensor): + tensor.zero_() + tokens = 3 * width + q = torch.ones((tokens, 2, 32), dtype=torch.bfloat16, device="cuda") + gates = torch.zeros((tokens, 2), dtype=torch.bfloat16, device="cuda") + a_log = torch.zeros(2, device="cuda") + bias = torch.zeros(2, device="cuda") + cu = torch.arange(4, device="cuda", dtype=torch.int32) * width if width > 1 else None + if cu is not None: + q = q.unsqueeze(0) + inputs = dict( + ssm_backend=ssm_backend, mode="replay", q=q, k=q, v=q, a=gates, b=gates, a_log=a_log, bias=bias, cu_seqlens=cu + ) + before = {name: x.clone() for name, x in vars(ssm_backend).items() if isinstance(x, torch.Tensor)} + + def oom(*args, **kwargs): + raise torch.cuda.OutOfMemoryError("scratch") + + tuner = tuning.select_config + monkeypatch.setattr(tuner, "rebuild_input_func", oom) + monkeypatch.setattr(tuner, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(tuner, "cached_configs", {}) + monkeypatch.setattr(tuner, "fast_match_configs", {}) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: 1) + monkeypatch.setattr( + tuner, "_bench", lambda *args, **kwargs: pytest.fail("Failed rebuild must not reach benchmarks") + ) + with pytest.raises(torch.cuda.OutOfMemoryError, match="scratch"): + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + tuning.resolve_run_config(**inputs) + assert not list(tmp_path.iterdir()) + config = {"BV": 32, "num_warps": 1, "num_stages": 3} + for name, expected in before.items(): + torch.testing.assert_close(getattr(ssm_backend, name), expected, atol=0, rtol=0) + + reqs = torch.arange(3, device="cuda", dtype=torch.int32) + ssm_backend.prepare_decode(reqs, cu) + tuning.replay_layer(ssm_backend, 0, q, q, q, gates, gates, a_log, bias, reqs, cu, run_config=config) + torch.cuda.synchronize() diff --git a/unit_tests/common/basemodel/triton_kernel/test_gated_rmsnorm_strided_gate.py b/unit_tests/common/basemodel/triton_kernel/test_gated_rmsnorm_strided_gate.py new file mode 100644 index 0000000000..64215880e1 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_gated_rmsnorm_strided_gate.py @@ -0,0 +1,41 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.norm.gated_rmsnorm import gated_rmsnorm_forward + + +@pytest.mark.parametrize("norm_before_gate", [True, False]) +def test_gated_rmsnorm_accepts_strided_3d_gate(norm_before_gate): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for gated RMSNorm test") + + torch.manual_seed(123) + tokens, heads, head_dim = 5, 4, 128 + packed_dim = heads * head_dim + 256 + x = torch.randn((tokens * heads, head_dim), device="cuda", dtype=torch.bfloat16) + weight = torch.randn((head_dim,), device="cuda", dtype=torch.bfloat16) + packed = torch.randn((tokens, packed_dim), device="cuda", dtype=torch.bfloat16) + z = packed[:, 128 : 128 + heads * head_dim].view(tokens, heads, head_dim) + assert not z.is_contiguous() + + run_config = {"BLOCK_N": head_dim, "num_warps": 1} + expected = gated_rmsnorm_forward( + x=x, + weight=weight, + bias=None, + eps=1e-6, + z=z.contiguous().view(-1, head_dim), + norm_before_gate=norm_before_gate, + run_config=run_config, + ) + actual = gated_rmsnorm_forward( + x=x, + weight=weight, + bias=None, + eps=1e-6, + z=z, + norm_before_gate=norm_before_gate, + run_config=run_config, + ) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/unit_tests/common/basemodel/triton_kernel/test_scalar_linear.py b/unit_tests/common/basemodel/triton_kernel/test_scalar_linear.py new file mode 100644 index 0000000000..2cbd3559d6 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_scalar_linear.py @@ -0,0 +1,70 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.scalar_linear import scalar_linear +from lightllm.common.quantization.scalar_projection import StableScalarProjection +from lightllm.common.quantization.quantize_method import WeightPack + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("hidden", [65, 2560]) +def test_reference_batching_strides_and_graph(dtype, hidden): + torch.manual_seed(2319) + input = torch.randn((336, hidden * 2), device="cuda", dtype=dtype)[1:335, ::2] + weight = torch.randn((1, hidden * 2), device="cuda", dtype=dtype)[:, ::2] + out = torch.empty((334, 3), device="cuda", dtype=dtype)[:, 1:2] + expected = (input.double() * weight.double()).sum(-1, keepdim=True).to(dtype) + assert scalar_linear(input, weight, out=out) is out + torch.testing.assert_close(out, expected, atol=0, rtol=0) + for size in (1, 3, 8, 24, 64, 128, 256, 334): + pieces = [scalar_linear(chunk, weight) for chunk in input.split(size)] + assert torch.equal(torch.cat(pieces), out) + permutation = torch.randperm(334, device="cuda") + assert torch.equal(scalar_linear(input[permutation], weight), out[permutation]) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + scalar_linear(input, weight, out=out) + input.mul_(0.5) + graph.replay() + torch.cuda.synchronize() + expected = (input.double() * weight.double()).sum(-1, keepdim=True).to(dtype) + torch.testing.assert_close(out, expected, atol=0, rtol=0) + + +def test_small_terms_survive_before_final_dtype_cast(): + input = torch.tensor([[1, 2 ** -8, 2 ** -24, 2 ** -24]], device="cuda", dtype=torch.bfloat16) + weight = torch.ones_like(input) + actual = scalar_linear(input, weight) + torch.testing.assert_close(actual, torch.full_like(actual, 1.0078125), atol=0, rtol=0) + + +def test_dispatch_preserves_weight_layout_output_and_fallbacks(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) + method = StableScalarProjection() + input = torch.randn((7, 65), device="cuda", dtype=torch.bfloat16) + weight = torch.randn((1, 65), device="cuda", dtype=input.dtype) + pack = WeightPack(weight=weight, weight_scale=None, weight_zero_point=None) + out = torch.empty((7, 1), device="cuda", dtype=input.dtype) + assert method.method_name == "none" + assert method.apply(input, pack, out=out, use_custom_tensor_mananger=False) is out + expected = (input.double() @ weight.double().t()).to(input.dtype) + torch.testing.assert_close(out, expected, atol=0, rtol=0) + torch.testing.assert_close(method.apply(input, pack), expected, atol=0, rtol=0) + + for outputs, bias in ((2, None), (1, torch.ones((1,), device="cuda", dtype=input.dtype))): + weight = torch.randn((outputs, 65), device="cuda", dtype=input.dtype) + pack = WeightPack(weight=weight, weight_scale=None, weight_zero_point=None) + actual = method.apply(input, pack, bias=bias, use_custom_tensor_mananger=False) + expected = torch.mm(input, weight.t()) if bias is None else torch.addmm(bias, input, weight.t()) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def test_empty_input_and_invalid_output(): + input = torch.empty((0, 65), device="cuda", dtype=torch.bfloat16) + weight = torch.ones((1, 65), device="cuda", dtype=input.dtype) + assert scalar_linear(input, weight).shape == (0, 1) + with pytest.raises(ValueError, match="output must match"): + scalar_linear(input, weight, out=torch.empty((0, 2), device="cuda", dtype=input.dtype)) diff --git a/unit_tests/common/gated_residual/test_gated_residual.py b/unit_tests/common/gated_residual/test_gated_residual.py new file mode 100644 index 0000000000..4833299460 --- /dev/null +++ b/unit_tests/common/gated_residual/test_gated_residual.py @@ -0,0 +1,178 @@ +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.gated_residual import GatedResidual + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def make_component(dim, branches, dtype, combine=True): + width, rank = dim * branches, 32 + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=dtype) + + return GatedResidual( + rand(width) * 0.1, + rand(rank, width) / width ** 0.5, + rand(width, rank) / rank ** 0.5, + rand(branches, width) / width ** 0.5 if combine else None, + branches, + ) + + +def reference_read(component, x): + grouped = x.float().unflatten(-1, (component.branches, -1)) + normalized = grouped * torch.rsqrt(grouped.square().mean(-1, keepdim=True) + component.eps) + normalized = (normalized.flatten(-2) * (1 + component.norm_weight.float())).to(x.dtype) + lowrank = F.silu(F.linear(normalized, component.down_weight) / component.branches) + gates = torch.sigmoid(F.linear(lowrank, component.up_weight)) + mixed = (normalized * gates).unflatten(-1, (component.branches, -1)).mean(-2) + if component.inject_weight is None: + return mixed + injection = 2 * torch.sigmoid(F.linear(normalized, component.inject_weight) / component.branches) + return mixed, injection + + +def reference_write(residual, update, gates): + return residual + (update.unsqueeze(-2) * gates.unsqueeze(-1)).flatten(-2) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("tokens,dim,branches", [(1, 65, 3), (17, 128, 4), (2, 2560, 4)]) +@pytest.mark.parametrize("combine", [True, False]) +def test_reference_parity(dtype, tokens, dim, branches, combine): + torch.manual_seed(37) + component = make_component(dim, branches, dtype, combine) + residual = torch.randn(1, tokens, dim * branches, device="cuda", dtype=dtype) + original = residual.clone() + actual, expected = component.read(residual), reference_read(component, residual) + + tolerance = max(1e-6, torch.finfo(dtype).eps) + if combine: + for result, ref in zip(actual, expected): + torch.testing.assert_close(result, ref, atol=tolerance, rtol=tolerance) + update = torch.randn(1, tokens, dim, device="cuda", dtype=dtype) + written = component.write(residual, update, actual[1]) + + torch.testing.assert_close(written, reference_write(residual, update, actual[1]), atol=0, rtol=0) + else: + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + torch.testing.assert_close(residual, original, atol=0, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("dim,branches", [(65, 3), (2560, 4)]) +def test_injection_gate_projection_is_independent_of_batch(dtype, dim, branches): + from lightllm.common.gated_residual.kernels import injection_gates + + torch.manual_seed(3107) + x = torch.randn(9, dim * branches, device="cuda", dtype=dtype) + weight = torch.randn(branches, dim * branches, device="cuda", dtype=dtype) * 0.01 + actual = injection_gates(x, weight, branches) + singles = torch.cat([injection_gates(row[None], weight, branches) for row in x]) + chunks = torch.cat([injection_gates(x[start : start + 3], weight, branches) for start in range(0, len(x), 3)]) + torch.testing.assert_close(actual, singles, atol=0, rtol=0) + torch.testing.assert_close(actual, chunks, atol=0, rtol=0) + expected = 2 * torch.sigmoid((x.double() @ weight.double().T).to(dtype) / branches) + tolerance = max(1e-6, torch.finfo(dtype).eps) + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + + +def test_saturated_mix_and_write(): + from lightllm.common.gated_residual.kernels import mix_streams + + component = make_component(65, 3, torch.bfloat16) + x = torch.randn(2, 195, device="cuda", dtype=torch.bfloat16) + logits = torch.linspace(-100, 100, 195, device="cuda", dtype=x.dtype).expand(2, -1).contiguous() + expected = (x * torch.sigmoid(logits)).view(2, 3, 65).mean(1) + torch.testing.assert_close(mix_streams(x, logits, 3), expected, atol=0, rtol=0) + gates = torch.tensor([[0, 1, 2], [2, 0, 1]], device="cuda", dtype=x.dtype) + update = torch.randn(2, 65, device="cuda", dtype=x.dtype) + torch.testing.assert_close(component.write(x, update, gates), reference_write(x, update, gates), atol=0, rtol=0) + + +def test_checkpoint_binding_and_graph_replay(): + component = make_component(128, 4, torch.bfloat16) + prefix = "model.language_model.layers.0.attn_hyper_connection" + weights = { + f"{prefix}.hc_norm.weight": component.norm_weight, + f"{prefix}.input_mix_weight_down.weight": component.down_weight, + f"{prefix}.input_mix_weight_up.weight": component.up_weight, + f"{prefix}.block_inject_weight.weight": component.inject_weight, + } + bound = GatedResidual.from_checkpoint(weights, prefix, branches=4) + assert bound.down_weight is component.down_weight + final = GatedResidual.from_checkpoint(weights, prefix, branches=4, use_combine=False) + residual = torch.randn(3, 512, device="cuda", dtype=torch.bfloat16) + update = torch.randn(3, 128, device="cuda", dtype=residual.dtype) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + _, gates = bound.read(residual) + bound.write(residual, update, gates) + final.read(residual) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + mixed, gates = bound.read(residual) + output = bound.write(residual, update, gates) + final_output = final.read(output) + for _ in range(2): + residual.normal_() + graph.replay() + expected_mixed, expected_gates = reference_read(bound, residual) + expected_output = reference_write(residual, update, expected_gates) + torch.testing.assert_close(mixed, expected_mixed, atol=2e-3, rtol=2e-3) + torch.testing.assert_close(output, expected_output, atol=2e-3, rtol=2e-3) + torch.testing.assert_close(final_output, reference_read(final, expected_output), atol=2e-3, rtol=2e-3) + + +def test_empty_and_invalid_inputs(): + component = make_component(65, 3, torch.bfloat16) + x = torch.empty(0, 195, device="cuda", dtype=torch.bfloat16) + mixed, gates = component.read(x) + assert component.write(x, mixed, gates).shape == x.shape + with pytest.raises(ValueError, match="last dimension"): + component.read(torch.empty(2, 194, device="cuda", dtype=x.dtype)) + with pytest.raises(ValueError, match="contiguous"): + component.read(torch.empty(2, 390, device="cuda", dtype=x.dtype)[:, ::2]) + with pytest.raises(ValueError, match="matching token"): + component.write(x, mixed, torch.empty(1, 3, device="cuda", dtype=x.dtype)) + with pytest.raises(ValueError, match="share"): + component.read(x.float()) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("dim,branches,rank", [(65, 3, 9), (2560, 4, 320)]) +def test_mixer_is_batch_invariant_and_matches_fp64(dtype, dim, branches, rank): + from lightllm.common.gated_residual.kernels import grouped_norm + + torch.manual_seed(701) + width = dim * branches + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=dtype) + + component = GatedResidual( + rand(width) * 0.1, + rand(rank, width) / width ** 0.5, + rand(width, rank) / rank ** 0.5, + rand(branches, width) / width ** 0.5, + branches, + ) + residual = rand(129, width) + expected = [component.read(row[None, :]) for row in residual] + scalar = tuple(torch.cat([item[index] for item in expected]) for index in range(2)) + for count in (3, 8, 16, 24, 64, 101, 128, 129): + outputs = [component.read(residual[start : start + count]) for start in range(0, 129, count)] + for index in range(2): + torch.testing.assert_close(torch.cat([item[index] for item in outputs]), scalar[index], atol=0, rtol=0) + normalized = grouped_norm(residual, component.norm_weight, branches, component.eps) + projected = F.linear(normalized.double(), component.down_weight.double()).to(dtype) + lowrank = F.silu(projected / branches) + logits = F.linear(lowrank.double(), component.up_weight.double()).to(dtype) + oracle = (normalized * torch.sigmoid(logits)).unflatten(-1, (branches, dim)).mean(-2) + torch.testing.assert_close(scalar[0], oracle, atol=2e-3, rtol=torch.finfo(dtype).eps) diff --git a/unit_tests/common/ngram_embedding/test_hash.py b/unit_tests/common/ngram_embedding/test_hash.py new file mode 100644 index 0000000000..ad217acf17 --- /dev/null +++ b/unit_tests/common/ngram_embedding/test_hash.py @@ -0,0 +1,206 @@ +import pytest +import torch + +from lightllm.common.ngram_embedding import NGramHasher + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +EOS = 248044 + + +def make_hasher(ngram_size=3, heads_per_ngram=8, overflow=False): + heads = (ngram_size - 1) * heads_per_ngram + multipliers = [39719048695937, 29866592791859, 5701618528319, 123456789123, 9876543212345][:ngram_size] + if overflow: + multipliers = [(1 << 62) + i * 141 + 1 for i in range(ngram_size)] + sizes = torch.arange(heads, device="cuda", dtype=torch.int64) * 2 + 20000003 + offsets = torch.cat([torch.zeros(1, device="cuda", dtype=torch.int64), sizes.cumsum(0)[:-1]]) + return NGramHasher(torch.tensor(multipliers, device="cuda", dtype=torch.int64), sizes, offsets, EOS) + + +def pack(sequences, dtype=torch.int64): + lengths = torch.tensor([0] + [len(seq) for seq in sequences], device="cuda", dtype=torch.int32) + tokens = torch.tensor([token for seq in sequences for token in seq], device="cuda", dtype=dtype) + return tokens, lengths.cumsum(0).to(torch.int32) + + +def reference(hasher, sequences, history): + rows, next_histories = [], [] + context = hasher.context_len + multipliers = hasher.multipliers.cpu() + sizes = hasher.head_vocab_sizes.tolist() + offsets = hasher.head_offsets.tolist() + per_order = hasher.heads // context + for request, sequence in enumerate(sequences): + full = history[request].tolist() + sequence + for index in range(context, len(full)): + shifted = [full[index]] + for shift in range(1, context + 1): + preceding = full[index - shift : index] + shifted.append(EOS if EOS in preceding else full[index - shift]) + row = [] + for order in range(2, context + 2): + mixed = torch.tensor(shifted[0], dtype=torch.int64) * multipliers[0] + for pos in range(1, order): + mixed ^= torch.tensor(shifted[pos], dtype=torch.int64) * multipliers[pos] + begin = (order - 2) * per_order + for head in range(begin, begin + per_order): + row.append((int(mixed) % sizes[head]) + offsets[head]) + rows.append(row) + next_histories.append(full[-context:]) + indices = torch.tensor(rows, device="cuda", dtype=torch.int64).reshape(-1, hasher.heads) + next_history = torch.tensor(next_histories, device="cuda", dtype=torch.int64).reshape(-1, context) + return indices, next_history + + +@pytest.mark.parametrize("ngram_size,heads_per_ngram", [(2, 1), (3, 8), (4, 3), (5, 2)]) +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("overflow", [False, True]) +def test_hash_matches_eos_segment_reference(ngram_size, heads_per_ngram, dtype, overflow): + hasher = make_hasher(ngram_size, heads_per_ngram, overflow) + sequences = [[], [7], [EOS], [11, EOS, 13, 14, EOS, EOS, 15, 16, 17], [], list(range(130))] + history = hasher.empty_history(len(sequences)) + history[0, -1] = 777 + history[3, -1] = 999 + if hasher.context_len > 1: + history[3, -2] = EOS + before = history.clone() + tokens, starts = pack(sequences, dtype) + actual = hasher.hash(tokens, starts, history) + expected = reference(hasher, sequences, history) + for result, ref in zip(actual, expected): + torch.testing.assert_close(result, ref, atol=0, rtol=0) + torch.testing.assert_close(history, before, atol=0, rtol=0) + assert torch.all(actual[0] >= hasher.head_offsets) + assert torch.all(actual[0] < hasher.head_offsets + hasher.head_vocab_sizes) + + +def test_chunked_prefill_and_decode_match_full_sequence(): + hasher = make_hasher() + sequences = [[1, 2, EOS, 3, 4, 5, EOS, 6], [11, EOS, 12, 13, EOS, 14], [EOS, 31]] + original = hasher.empty_history(len(sequences)) + expected, full_history = hasher.hash(*pack(sequences), original) + history = original + per_request = [[] for _ in sequences] + consumed = [0] * len(sequences) + for counts in [(1, 0, 1), (2, 3, 1), (1, 1, 0), (4, 2, 0)]: + chunks = [sequence[start : start + count] for sequence, start, count in zip(sequences, consumed, counts)] + tokens, starts = pack(chunks) + result, history = hasher.hash(tokens, starts, history) + for request, count in enumerate(counts): + per_request[request].append(result[int(starts[request]) : int(starts[request + 1])]) + consumed[request] += count + combined = torch.cat([torch.cat(chunks) for chunks in per_request]) + torch.testing.assert_close(combined, expected, atol=0, rtol=0) + torch.testing.assert_close(history, full_history, atol=0, rtol=0) + + +def test_speculative_rejection_and_cache_snapshot(): + hasher = make_hasher() + sequences = [[1, 2, 3, EOS], [11, EOS, 12], [21, 22], []] + tokens, starts = pack(sequences) + initial = hasher.empty_history(len(sequences)) + initial[:, -1] = 99 + snapshot = initial.clone() + + hasher.hash(tokens, starts, initial) + accepted = torch.tensor([1, 2, 0, 0], device="cuda", dtype=torch.int32) + committed = hasher.history_after(tokens, starts, initial, accepted) + expected = reference(hasher, [seq[:n] for seq, n in zip(sequences, [1, 2, 0, 0])], snapshot)[1] + torch.testing.assert_close(committed, expected, atol=0, rtol=0) + torch.testing.assert_close(initial, snapshot, atol=0, rtol=0) + + followup = [[7], [8], [9], [10]] + result = hasher.hash(*pack(followup), committed) + ref = reference(hasher, followup, expected) + for actual, expected in zip(result, ref): + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + order = torch.tensor([2, 0, 2], device="cuda") + histories = snapshot.index_select(0, order) + sequences = [[31], [32, EOS, 33], [34]] + for actual, expected in zip(hasher.hash(*pack(sequences), histories), reference(hasher, sequences, histories)): + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def test_checkpoint_and_cuda_graph_replay(): + source = make_hasher() + prefix = "model.language_model.layers.1.ple.ple_embedding" + buffers = { + f"{prefix}.layer_multipliers": source.multipliers, + f"{prefix}.ngram_heads_vocab_sizes": source.head_vocab_sizes, + f"{prefix}.ngram_heads_offsets": source.head_offsets, + } + hasher = NGramHasher.from_checkpoint(buffers, prefix, EOS) + assert hasher.multipliers is source.multipliers + history = hasher.empty_history(3) + tokens, starts = pack([[1, 2, 3], [11, 12], [21, 22]]) + accepted = torch.tensor([3, 1, 2], dtype=torch.int32, device="cuda") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + hasher.hash(tokens, starts, history) + hasher.history_after(tokens, starts, history, accepted) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + indices, next_history = hasher.hash(tokens, starts, history) + committed = hasher.history_after(tokens, starts, history, accepted) + for sequences, counts in [ + ([[EOS, 2], [], [11, 12, 13, 14, 15]], [1, 0, 4]), + ([[], [31, EOS, 32], [41, 42, 43, 44]], [0, 2, 0]), + ]: + updated_tokens, updated_starts = pack(sequences) + tokens.copy_(updated_tokens) + starts.copy_(updated_starts) + accepted.copy_(torch.tensor(counts, device="cuda", dtype=torch.int32)) + history[:, -1] = 151 + graph.replay() + ref_indices, ref_history = reference(hasher, sequences, history) + torch.testing.assert_close(indices, ref_indices, atol=0, rtol=0) + torch.testing.assert_close(next_history, ref_history, atol=0, rtol=0) + ref_committed = reference(hasher, [seq[:n] for seq, n in zip(sequences, counts)], history)[1] + torch.testing.assert_close(committed, ref_committed, atol=0, rtol=0) + + +def test_empty_batches_and_bad_layouts(): + hasher = make_hasher() + for sequences in [[], [[], [], []]]: + history = hasher.empty_history(len(sequences)) + indices, next_history = hasher.hash(*pack(sequences), history) + assert indices.shape == (0, hasher.heads) + torch.testing.assert_close(next_history, history, atol=0, rtol=0) + tokens, starts = pack([[1, 2, 3]]) + history = hasher.empty_history(1) + with pytest.raises(ValueError, match="packed"): + hasher.hash(tokens.view(1, -1), starts, history) + with pytest.raises(ValueError, match="Starts"): + hasher.hash(tokens, starts.long(), history) + with pytest.raises(ValueError, match="History"): + hasher.hash(tokens, starts, history.int()) + with pytest.raises(ValueError, match="contiguous"): + hasher.hash(tokens, starts, torch.empty(1, 4, device="cuda", dtype=torch.int64)[:, ::2]) + with pytest.raises(ValueError, match="Accepted"): + hasher.history_after(tokens, starts, history, torch.ones(1, device="cuda", dtype=torch.int64)) + + +@pytest.mark.parametrize("invalid", ["multipliers", "empty_heads", "dtype", "zero_size", "overlap"]) +def test_invalid_checkpoint_buffers(invalid): + from dataclasses import replace + + hasher = make_hasher() + with pytest.raises(ValueError): + if invalid == "multipliers": + replace(hasher, multipliers=hasher.multipliers[:1]) + elif invalid == "empty_heads": + replace(hasher, head_vocab_sizes=hasher.head_vocab_sizes[:0]) + elif invalid == "dtype": + replace(hasher, multipliers=hasher.multipliers.int()) + elif invalid == "zero_size": + sizes = hasher.head_vocab_sizes.clone() + sizes[0] = 0 + replace(hasher, head_vocab_sizes=sizes) + else: + offsets = hasher.head_offsets.clone() + offsets[1] = 0 + replace(hasher, head_offsets=offsets) diff --git a/unit_tests/common/ngram_embedding/test_pinned_embedding.py b/unit_tests/common/ngram_embedding/test_pinned_embedding.py new file mode 100644 index 0000000000..57baa1744f --- /dev/null +++ b/unit_tests/common/ngram_embedding/test_pinned_embedding.py @@ -0,0 +1,88 @@ +import importlib +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights.pinned_embedding_weight import PinnedRowShardedEmbeddingWeight +from lightllm.models.qwen4_exp import pinned_memory + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.fixture(autouse=True) +def runtime(monkeypatch): + base = importlib.import_module("lightllm.common.basemodel.layer_weights.meta_weights.base_weight") + platform = importlib.import_module("lightllm.common.basemodel.layer_weights.meta_weights.platform_op") + monkeypatch.setattr(base, "get_current_device_id", torch.cuda.current_device) + monkeypatch.setattr( + platform, + "get_env_start_args", + lambda: SimpleNamespace(hardware_platform="cuda", enable_torch_fallback=False, enable_triton_fallback=False), + ) + monkeypatch.setattr(pinned_memory, "_LARGE_BYTES", 1) + monkeypatch.setattr(pinned_memory, "_CHUNK_BYTES", 65536) + + +def weight(dtype, rank=0, world=1): + return PinnedRowShardedEmbeddingWeight( + 160, 1024, "table.weight", dtype, 4, cpu_allocator=pinned_memory.pinned_zeros, tp_rank=rank, tp_world_size=world + ) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_streamed_host_shards_and_vocab_parallel_gather(dtype): + table = (torch.arange(1024 * 160).reshape(1024, 160) % 127 - 64).to(dtype) + ids = torch.tensor([0, 255, 256, 340, 341, 511, 512, 681, 682, 1023, -1, 1024], device="cuda") + outputs = [] + for rank in range(3): + value = weight(dtype, rank, 3) + address = value.weight.data_ptr() + assert not value.verify_load() and value.weight.device.type == "cpu" and value.weight.is_pinned() + assert len(value.weight._qwen4_pinned_allocation.addresses) > 1 + for index in (2, 0, 3, 1): + value.load_hf_weights({f"table.shard_{index}.weight": table[index * 256 : (index + 1) * 256]}) + assert value.verify_load() and value.weight.data_ptr() == address + torch.testing.assert_close(value.weight, table[value.tp_vocab_start_id : value.tp_vocab_end_id], atol=0, rtol=0) + output = value(ids) + assert output.is_cuda + outputs.append(output) + expected = table.cuda().index_select(0, ids.clamp(0, 1023)) + expected[(ids < 0) | (ids >= 1024)] = 0 + torch.testing.assert_close(torch.stack(outputs).sum(0), expected, atol=0, rtol=0) + + +def test_device_checkpoint_staging_and_dynamic_graph_replay(): + table = (torch.arange(1024 * 160).reshape(1024, 160) % 127).to(torch.bfloat16).cuda() + value = weight(torch.bfloat16) + value.load_hf_weights({"table.weight": table}) + assert value.verify_load() + ids = torch.tensor([0, 191, 192, 511, 512, 1023], device="cuda") + output = torch.empty((ids.numel(), 160), dtype=table.dtype, device="cuda") + value(ids, out=output) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + value(ids, out=output) + ids.copy_(torch.tensor([1023, 512, 511, 192, 191, 0], device="cuda")) + graph.replay() + torch.testing.assert_close(output, table.index_select(0, ids), atol=0, rtol=0) + graph.reset() + + +def test_rejects_pageable_allocator_and_wrong_output(): + with pytest.raises(ValueError, match="pinned CPU"): + PinnedRowShardedEmbeddingWeight( + 160, + 1024, + "table.weight", + torch.bfloat16, + 4, + cpu_allocator=lambda shape, **kwargs: torch.zeros(shape, dtype=kwargs["dtype"]), + tp_rank=0, + tp_world_size=1, + ) + value = weight(torch.bfloat16) + with pytest.raises(ValueError, match="CUDA vector"): + value(torch.tensor([0])) + with pytest.raises(ValueError, match="output"): + value(torch.tensor([0], device="cuda"), out=torch.empty((1, 160), device="cuda", dtype=torch.float32)) diff --git a/unit_tests/common/ngram_embedding/test_ple.py b/unit_tests/common/ngram_embedding/test_ple.py new file mode 100644 index 0000000000..677a923734 --- /dev/null +++ b/unit_tests/common/ngram_embedding/test_ple.py @@ -0,0 +1,259 @@ +import math +from dataclasses import replace + +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.ngram_embedding.ple import PLELayer + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def make_layer(dim=65, branches=3, embed_dim=32, kernel=4, dilation=3, dtype=torch.bfloat16): + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=dtype) + + width = dim * branches + return PLELayer( + rand(width, embed_dim) / embed_dim ** 0.5, + rand(dim, embed_dim) / embed_dim ** 0.5, + rand(width) * 0.1, + rand(width) * 0.1, + rand(width) * 0.1, + rand(width, kernel) * 0.1, + branches, + dilation, + ) + + +def norm(x, weight, branches, eps): + grouped = x.float().unflatten(-1, (branches, -1)) + grouped = grouped * torch.rsqrt(grouped.square().mean(-1, keepdim=True) + eps) + return (grouped.flatten(-2) * (1 + weight.float())).to(x.dtype) + + +def reference_values(layer, residual, embeddings, mask=None): + key = norm(F.linear(embeddings, layer.key_weight), layer.norm_key, layer.branches, layer.eps) + query = norm(residual, layer.norm_query, layer.branches, layer.eps) + key, query = [x.unflatten(-1, (layer.branches, layer.hidden_size)) for x in [key, query]] + value = F.linear(embeddings, layer.value_weight) + gate = (key * query).sum(-1, keepdim=True) / math.sqrt(layer.hidden_size) + gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() + gated = (torch.sigmoid(gate) * value.unsqueeze(-2)).flatten(-2) + normalized = norm(gated, layer.norm_conv, layer.branches, layer.eps) + if mask is not None: + gated = gated * mask[:, None] + normalized = normalized * mask[:, None] + return gated, normalized + + +def reference_conv(layer, values, gated, lengths, history): + result, next_state = [], [] + offset = 0 + for request, count in enumerate(lengths): + current = values[offset : offset + count] + full = torch.cat([history[request], current]) + if count: + conv = ( + F.conv1d( + full.T.unsqueeze(0), + layer.conv_weight.unsqueeze(1), + dilation=layer.dilation, + groups=values.shape[1], + ) + .squeeze(0) + .T + ) + result.append(gated[offset : offset + count] + F.silu(conv)) + next_state.append(full[-layer.context_len :] if layer.context_len else full[:0]) + offset += count + output = torch.cat(result) if result else values.new_empty((0, values.shape[1])) + state = torch.stack(next_state) if next_state else history.clone() + return output, state + + +def starts_for(lengths): + return torch.tensor([0] + list(lengths), device="cuda", dtype=torch.int32).cumsum(0).to(torch.int32) + + +def assert_close(actual, expected, dtype): + tolerance = max(2e-6, torch.finfo(dtype).eps * 2) + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("dim,branches,kernel,dilation", [(65, 3, 4, 3), (128, 4, 3, 2), (32, 2, 1, 1)]) +@pytest.mark.parametrize("masked", [False, True]) +def test_ple_reference_equations(dtype, dim, branches, kernel, dilation, masked): + torch.manual_seed(137) + layer = make_layer(dim, branches, kernel=kernel, dilation=dilation, dtype=dtype) + lengths = [0, 1, 2, 13, 0, 7] + count = sum(lengths) + residual = torch.randn(count, dim * branches, device="cuda", dtype=dtype) + embeddings = torch.randn(count, 32, device="cuda", dtype=dtype) + mask = (torch.arange(count, device="cuda") % 3 != 0) if masked else None + history = layer.empty_history(len(lengths)).normal_() + original = history.clone() + actual, conv_inputs = layer.forward(residual, embeddings, starts_for(lengths), history, mask) + gated, expected_conv_inputs = reference_values(layer, residual, embeddings, mask) + expected, expected_state = reference_conv(layer, expected_conv_inputs, gated, lengths, history) + assert_close(conv_inputs, expected_conv_inputs, dtype) + assert_close(actual, expected, dtype) + assert_close(layer.history_after(conv_inputs, starts_for(lengths), history), expected_state, dtype) + torch.testing.assert_close(history, original, atol=0, rtol=0) + + +def test_production_hidden_width(): + torch.manual_seed(137) + layer = make_layer(dim=2560, branches=4, embed_dim=2560) + residual = torch.randn(2, 10240, device="cuda", dtype=torch.bfloat16) + embeddings = torch.randn(2, 2560, device="cuda", dtype=torch.bfloat16) + history = layer.empty_history(2) + output, inputs = layer.forward(residual, embeddings, starts_for([1, 1]), history) + gated, expected_inputs = reference_values(layer, residual, embeddings) + expected, _ = reference_conv(layer, expected_inputs, gated, [1, 1], history) + assert_close(inputs, expected_inputs, residual.dtype) + assert_close(output, expected, residual.dtype) + + +def test_chunking_speculative_acceptance_and_request_isolation(): + torch.manual_seed(137) + layer = make_layer(dim=32, branches=2, dtype=torch.float32) + lengths = [8, 6, 2] + starts = starts_for(lengths) + residual = torch.randn(sum(lengths), 64, device="cuda") + embeddings = torch.randn(sum(lengths), 32, device="cuda") + history = layer.empty_history(3).normal_() + full, full_inputs = layer.forward(residual, embeddings, starts, history) + full_state = layer.history_after(full_inputs, starts, history) + state = history.clone() + consumed = [0] * 3 + per_request = [[] for _ in range(3)] + for counts in [(1, 0, 1), (2, 3, 1), (1, 1, 0), (4, 2, 0)]: + slices = [slice(int(starts[i]) + consumed[i], int(starts[i]) + consumed[i] + n) for i, n in enumerate(counts)] + chunks = [torch.cat([tensor[s] for s in slices]) for tensor in [residual, embeddings]] + chunk_starts = starts_for(counts) + output, conv_inputs = layer.forward(*chunks, chunk_starts, state) + state = layer.history_after(conv_inputs, chunk_starts, state) + for i, n in enumerate(counts): + per_request[i].append(output[int(chunk_starts[i]) : int(chunk_starts[i + 1])]) + consumed[i] += n + assert_close(torch.cat([torch.cat(parts) for parts in per_request]), full, residual.dtype) + assert_close(state, full_state, residual.dtype) + accepted = torch.tensor([1, 4, 0], device="cuda", dtype=torch.int32) + committed = layer.history_after(full_inputs, starts, history, accepted) + expected = torch.stack( + [ + torch.cat([history[i], full_inputs[int(starts[i]) : int(starts[i]) + n]])[-layer.context_len :] + for i, n in enumerate([1, 4, 0]) + ] + ) + torch.testing.assert_close(committed, expected, atol=0, rtol=0) + + order = torch.tensor([2, 0, 2], device="cuda") + reordered = committed.index_select(0, order) + new_residual, new_embeddings = residual[:3], embeddings[:3] + actual, _ = layer.forward(new_residual, new_embeddings, starts_for([1, 1, 1]), reordered) + gated, normalized = reference_values(layer, new_residual, new_embeddings) + expected, _ = reference_conv(layer, normalized, gated, [1, 1, 1], reordered) + assert_close(actual, expected, residual.dtype) + + +def test_checkpoint_binding_and_graph_replay(): + layer = make_layer(dim=32, branches=2) + prefix = "model.language_model.layers.1.ple" + weights = { + f"{prefix}.{name}": tensor + for name, tensor in [ + ("key_proj.weight", layer.key_weight), + ("value_proj.weight", layer.value_weight), + ("norm_key.weight", layer.norm_key), + ("norm_query.weight", layer.norm_query), + ("norm_conv.weight", layer.norm_conv), + ("conv1d.weight", layer.conv_weight.unsqueeze(1)), + ] + } + bound = PLELayer.from_checkpoint(weights, prefix, branches=2, dilation=3) + assert bound.key_weight is layer.key_weight + assert bound.conv_weight.data_ptr() == layer.conv_weight.data_ptr() + residual = torch.randn(7, 64, device="cuda", dtype=torch.bfloat16) + embeddings = torch.randn(7, 32, device="cuda", dtype=torch.bfloat16) + starts = starts_for([3, 2, 2]) + accepted = torch.tensor([3, 1, 2], device="cuda", dtype=torch.int32) + mask = torch.ones(7, device="cuda", dtype=torch.bool) + history = layer.empty_history(3) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + _, conv_inputs = bound.forward(residual, embeddings, starts, history, mask) + bound.history_after(conv_inputs, starts, history, accepted) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output, conv_inputs = bound.forward(residual, embeddings, starts, history, mask) + next_state = bound.history_after(conv_inputs, starts, history, accepted) + for lengths, counts in [([2, 0, 5], [1, 0, 4]), ([0, 3, 4], [0, 2, 0])]: + residual.normal_() + embeddings.normal_() + history.normal_() + starts.copy_(starts_for(lengths)) + accepted.copy_(torch.tensor(counts, device="cuda", dtype=torch.int32)) + mask[::2] = False + graph.replay() + gated, expected_inputs = reference_values(bound, residual, embeddings, mask) + expected, _ = reference_conv(bound, expected_inputs, gated, lengths, history) + assert_close(output, expected, residual.dtype) + expected_state = [] + for i, count in enumerate(counts): + expected_state.append( + torch.cat([history[i], expected_inputs[int(starts[i]) : int(starts[i]) + count]])[-bound.context_len :] + ) + assert_close(next_state, torch.stack(expected_state), residual.dtype) + + +def test_empty_batches_and_invalid_inputs(): + layer = make_layer(dim=32, branches=2) + for lengths in [[], [0, 0, 0]]: + residual = layer.key_weight.new_empty((0, 64)) + embeddings = layer.key_weight.new_empty((0, 32)) + history = layer.empty_history(len(lengths)) + output, conv_inputs = layer.forward(residual, embeddings, starts_for(lengths), history) + assert output.shape == (0, 64) + torch.testing.assert_close( + layer.history_after(conv_inputs, starts_for(lengths), history), history, atol=0, rtol=0 + ) + with pytest.raises(ValueError, match="projection"): + replace(layer, value_weight=layer.value_weight[:1]) + with pytest.raises(ValueError, match="convolution"): + replace(layer, conv_weight=layer.conv_weight[:, :0]) + residual = layer.key_weight.new_empty((1, 64)) + embeddings = layer.key_weight.new_empty((1, 32)) + with pytest.raises(ValueError, match="history"): + layer.forward(residual, embeddings, starts_for([1]), layer.empty_history(2)) + with pytest.raises(ValueError, match="mask"): + layer.forward(residual, embeddings, starts_for([1]), layer.empty_history(1), torch.ones(1, device="cuda")) + with pytest.raises(ValueError, match="accepted"): + layer.history_after(residual, starts_for([1]), layer.empty_history(1), torch.ones(1, device="cuda")) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_gate_zero_near_zero_and_signed_saturation(dtype): + from lightllm.common.ngram_embedding.ple_kernels import _gate + + dots = torch.tensor([0, 1e-8, -1e-8, 1e-5, -1e-5, 100, -100], device="cuda", dtype=dtype) + key = torch.zeros(7, 4, device="cuda", dtype=dtype) + query = torch.zeros_like(key) + key[:, 0] = dots * 2 + query[:, 0] = 1 + value = torch.tensor([[-2, -1, 1, 2]], device="cuda", dtype=dtype).expand(7, -1).contiguous() + gate = (key * query).sum(-1, keepdim=True) / 2 + gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() + expected = torch.sigmoid(gate) * value + actual = torch.empty_like(key) + _gate[(7, 1)](key, query, value, None, actual, 4, 1, False, 4, enable_fp_fusion=False) + if dtype == torch.float32: + torch.testing.assert_close(actual, expected, atol=5e-7, rtol=5e-7) + else: + torch.testing.assert_close(actual, expected, atol=0, rtol=0) diff --git a/unit_tests/common/ngram_embedding/test_row_sharded_embedding.py b/unit_tests/common/ngram_embedding/test_row_sharded_embedding.py new file mode 100644 index 0000000000..cacb940aa4 --- /dev/null +++ b/unit_tests/common/ngram_embedding/test_row_sharded_embedding.py @@ -0,0 +1,161 @@ +from types import SimpleNamespace +import importlib + +import pytest +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights.row_sharded_embedding_weight import RowShardedEmbeddingWeight + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +PREFIX = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding" + + +@pytest.fixture(autouse=True) +def runtime_configuration(monkeypatch): + base = importlib.import_module("lightllm.common.basemodel.layer_weights.meta_weights.base_weight") + platform = importlib.import_module("lightllm.common.basemodel.layer_weights.meta_weights.platform_op") + monkeypatch.setattr(base, "get_current_device_id", lambda: torch.cuda.current_device()) + monkeypatch.setattr( + platform, + "get_env_start_args", + lambda: SimpleNamespace( + hardware_platform="cuda", + enable_torch_fallback=False, + enable_triton_fallback=False, + ), + ) + + +def make_weight(rank=0, world=1): + return RowShardedEmbeddingWeight(5, 24, f"{PREFIX}.weight", torch.bfloat16, 4, rank, world) + + +def shards(table): + return [{f"{PREFIX}.shard_{i}.weight": part} for i, part in enumerate(table.chunk(4))] + + +@pytest.mark.parametrize("world", [1, 2, 3, 8]) +def test_streamed_shards_and_vocab_parallel_lookup(world): + table = torch.arange(120, dtype=torch.bfloat16).view(24, 5) + checkpoint = shards(table) + ids = torch.tensor([0, 5, 6, 7, 8, 11, 12, 15, 17, 18, 23], device="cuda", dtype=torch.int64) + outputs = [] + for rank in range(world): + weight = make_weight(rank, world) + assert not weight.verify_load() + required, loaded = set(), set() + for i in range(4): + if max(i * 6, weight.tp_vocab_start_id) < min((i + 1) * 6, weight.tp_vocab_end_id): + required.add(i) + for index in [2, 0, 3, 1]: + weight.load_hf_weights(checkpoint[index]) + loaded.add(index) + assert weight.verify_load() == (loaded >= required) + torch.testing.assert_close( + weight.weight.cpu(), table[weight.tp_vocab_start_id : weight.tp_vocab_end_id], atol=0, rtol=0 + ) + outputs.append(weight(ids)) + + torch.testing.assert_close( + torch.stack(outputs).sum(0), torch.nn.functional.embedding(ids, table.cuda()), atol=0, rtol=0 + ) + + +def test_missing_shards_bad_shapes_and_other_prefixes(): + weight = make_weight() + table = torch.arange(120, dtype=torch.bfloat16).view(24, 5) + checkpoint = shards(table) + for index in [0, 1, 3]: + weight.load_hf_weights(checkpoint[index]) + assert not weight.verify_load() + weight.load_hf_weights({"other.shard_2.weight": table[:6]}) + assert not weight.verify_load() + with pytest.raises(ValueError, match="shard 2"): + weight.load_hf_weights({f"{PREFIX}.shard_2.weight": table[:5]}) + assert not weight.verify_load() + weight.load_hf_weights(checkpoint[2]) + assert weight.verify_load() + + weight.load_hf_weights({"other.weight": table}) + assert weight.verify_load() + + +def test_full_checkpoint_compatibility_and_reject_mixed_formats(): + table = torch.arange(120, dtype=torch.bfloat16).view(24, 5) + weight = make_weight(1, 3) + weight.load_hf_weights({f"{PREFIX}.weight": table}) + assert weight.verify_load() + weight.load_hf_weights({"other.weight": table}) + assert weight.verify_load() + torch.testing.assert_close(weight.weight.cpu(), table[8:16], atol=0, rtol=0) + with pytest.raises(ValueError, match="mix"): + weight.load_hf_weights(shards(table)[0]) + partial = make_weight() + partial.load_hf_weights(shards(table)[1]) + with pytest.raises(ValueError, match="mix"): + partial.load_hf_weights({f"{PREFIX}.weight": table}) + with pytest.raises(ValueError, match="dimensions"): + make_weight().load_hf_weights({f"{PREFIX}.weight": table[:5]}) + + +def test_row_shard_configuration_and_graph_replay(): + with pytest.raises(ValueError, match="divisible"): + RowShardedEmbeddingWeight(5, 23, f"{PREFIX}.weight", torch.bfloat16, 4, 0, 1) + table = torch.arange(120, dtype=torch.bfloat16).view(24, 5) + weight = make_weight() + for part in shards(table): + weight.load_hf_weights(part) + ids = torch.tensor([0, 6, 12, 23], device="cuda", dtype=torch.int64) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + weight(ids) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = weight(ids) + for values in [[23, 12, 6, 0], [1, 7, 13, 19]]: + ids.copy_(torch.tensor(values, device="cuda")) + graph.replay() + torch.testing.assert_close(output.cpu(), table[values], atol=0, rtol=0) + + +def test_packed_hash_lookup_and_ple_composition(): + from lightllm.common.ngram_embedding import NGramHasher + from lightllm.common.ngram_embedding.ple import PLELayer + + torch.manual_seed(137) + sizes = torch.tensor([2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53], device="cuda", dtype=torch.int64) + offsets = torch.cat([sizes.new_zeros(1), sizes.cumsum(0)[:-1]]) + hasher = NGramHasher( + torch.tensor([23703573157769, 20109073645365, 8052911324071], device="cuda"), sizes, offsets, 248044 + ) + tokens = torch.tensor([1, 2, 248044, 3, 4, 5, 6, 7], device="cuda", dtype=torch.int64) + starts = torch.tensor([0, 3, 4, 4, 8], device="cuda", dtype=torch.int32) + indices, _ = hasher.hash(tokens, starts, hasher.empty_history(4)) + + table = torch.randn(512, 4, dtype=torch.bfloat16) + partition_outputs = [] + for rank in range(3): + weight = RowShardedEmbeddingWeight(4, 512, f"{PREFIX}.weight", torch.bfloat16, 128, rank, 3) + for index, shard in reversed(list(enumerate(table.chunk(128)))): + weight.load_hf_weights({f"{PREFIX}.shard_{index}.weight": shard}) + assert weight.verify_load() + partition_outputs.append(weight(indices.flatten())) + embeddings = torch.stack(partition_outputs).sum(0).view(tokens.numel(), 64) + expected_embeddings = torch.nn.functional.embedding(indices, table.cuda()).flatten(-2) + torch.testing.assert_close(embeddings, expected_embeddings, atol=0, rtol=0) + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) + + layer = PLELayer( + rand(64, 64) / 8, rand(32, 64) / 8, rand(64) * 0.1, rand(64) * 0.1, rand(64) * 0.1, rand(64, 4) * 0.1, 2, 3 + ) + history = layer.empty_history(4) + residual = rand(tokens.numel(), 64) + output, conv_inputs = layer.forward(residual, embeddings, starts, history) + expected_output, expected_inputs = layer.forward(residual, expected_embeddings, starts, history) + torch.testing.assert_close(output, expected_output, atol=0, rtol=0) + torch.testing.assert_close(conv_inputs, expected_inputs, atol=0, rtol=0) diff --git a/unit_tests/common/qsa/test_attention.py b/unit_tests/common/qsa/test_attention.py new file mode 100644 index 0000000000..e935056012 --- /dev/null +++ b/unit_tests/common/qsa/test_attention.py @@ -0,0 +1,133 @@ +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.qsa.attention import indexed_attention + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def reference(query, key, value, indices, scale=None): + selected_k = key[indices.clamp_min(0)].transpose(1, 2) + selected_v = value[indices.clamp_min(0)].transpose(1, 2) + mask = (indices >= 0)[:, None, None, :] + return F.scaled_dot_product_attention( + query.unsqueeze(2), selected_k, selected_v, attn_mask=mask, scale=scale, enable_gqa=True + ).squeeze(2) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize( + "rows,qheads,kvheads,dim,vdim,count", [(3, 6, 2, 32, 16, 79), (1, 24, 2, 256, 256, 2051), (5, 3, 1, 65, 33, 11)] +) +@pytest.mark.parametrize("splits", [1, 4, None]) +@pytest.mark.parametrize("bounded", [False, True]) +def test_sparse_gqa_matches_sdpa(dtype, rows, qheads, kvheads, dim, vdim, count, splits, bounded): + torch.manual_seed(437) + slots = count + 31 + q = torch.randn(rows, qheads, dim, device="cuda", dtype=dtype) + k = torch.randn(slots, kvheads, dim, device="cuda", dtype=dtype) + v = torch.randn(slots, kvheads, vdim, device="cuda", dtype=dtype) + indices = torch.stack([torch.randperm(slots, device="cuda")[:count] for _ in range(rows)]).int() + indices[:, ::7] = -1 + lengths = torch.full((rows,), count + 17, device="cuda", dtype=torch.int32) if bounded else None + if bounded: + indices[0, count // 2 :] = -1 + lengths[0] = count // 2 + if rows > 1: + indices[-1].fill_(-1) + if bounded: + lengths[-1] = -1 + before = (q.clone(), k.clone(), v.clone(), indices.clone()) + actual = indexed_attention(q, k, v, indices, splits=splits, selection_lengths=lengths) + expected = reference(q, k, v, indices) + tolerance = 2e-5 if dtype == torch.float32 else torch.finfo(dtype).eps * 2 + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + for tensor, original in zip((q, k, v, indices), before): + torch.testing.assert_close(tensor, original, atol=0, rtol=0) + if rows > 1: + assert torch.count_nonzero(actual[-1]) == 0 + + +@pytest.mark.parametrize("bounded", [False, True]) +def test_attention_cuda_graph_replay_and_custom_scale(bounded): + q = torch.randn(2, 6, 32, device="cuda", dtype=torch.bfloat16) + k = torch.randn(128, 2, 32, device="cuda", dtype=q.dtype) + v = torch.randn(128, 2, 16, device="cuda", dtype=q.dtype) + indices = torch.arange(79, device="cuda", dtype=torch.int32).expand(2, -1).contiguous() + lengths = torch.full((2,), 79, device="cuda", dtype=torch.int32) if bounded else None + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + indexed_attention(q, k, v, indices, scale=0.17, selection_lengths=lengths) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = indexed_attention(q, k, v, indices, scale=0.17, selection_lengths=lengths) + for bounds in ([13, 0], [128, 63]): + q.normal_() + indices.copy_(torch.stack([torch.randperm(128, device="cuda")[:79] for _ in range(2)]).int()) + indices[:, ::3] = -1 + if bounded: + lengths.copy_(torch.tensor(bounds, device="cuda", dtype=torch.int32)) + indices.masked_fill_(torch.arange(79, device="cuda")[None, :] >= lengths[:, None], -1) + graph.replay() + torch.testing.assert_close(output, reference(q, k, v, indices, scale=0.17), atol=0.016, rtol=0.016) + + +def test_empty_selection_and_invalid_shapes(): + q = torch.randn(2, 6, 32, device="cuda", dtype=torch.bfloat16) + k = torch.randn(10, 2, 32, device="cuda", dtype=q.dtype) + v = torch.randn(10, 2, 16, device="cuda", dtype=q.dtype) + ids = torch.empty(2, 0, device="cuda", dtype=torch.int32) + assert torch.count_nonzero(indexed_attention(q, k, v, ids)) == 0 + assert indexed_attention(q[:0], k, v, ids[:0]).shape == (0, 6, 16) + with pytest.raises(ValueError, match="divisible"): + indexed_attention(q[:, :5].contiguous(), k, v, ids) + with pytest.raises(ValueError, match="int32"): + indexed_attention(q, k, v, ids.long()) + with pytest.raises(ValueError, match="Selection lengths"): + indexed_attention(q, k, v, ids, selection_lengths=torch.zeros(2, device="cuda", dtype=torch.int64)) + + +def test_native_interleaved_kv_views_match_contiguous_cache(): + q = torch.randn(3, 24, 256, device="cuda", dtype=torch.bfloat16) + packed = torch.randn(97, 4, 256, device="cuda", dtype=q.dtype) + k, v = packed[:, :2], packed[:, 2:] + assert not k.is_contiguous() and not v.is_contiguous() + indices = torch.stack([torch.randperm(97, device="cuda")[:79] for _ in range(3)]).int() + indices[:, ::7] = -1 + actual = indexed_attention(q, k, v, indices) + expected = indexed_attention(q, k.contiguous(), v.contiguous(), indices) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_direct_short_rows_preserve_fixed_partition_and_graph_replay(dtype): + torch.manual_seed(701) + rows, slots, width = 5, 4096, 2051 + q = torch.randn(rows, 6, 256, device="cuda", dtype=dtype) + k = torch.randn(slots, 1, 256, device="cuda", dtype=dtype) + v = torch.randn_like(k) + ids = torch.stack([torch.randperm(slots, device="cuda")[:width] for _ in range(rows)]).int() + lengths = torch.tensor([0, 16, 64, 65, 2048], device="cuda", dtype=torch.int32) + + def compute(direct): + return indexed_attention(q, k, v, ids, splits=32, selection_lengths=lengths, short_rows_direct=direct) + + torch.testing.assert_close(compute(True), compute(False), atol=0, rtol=0) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + compute(True) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = compute(True) + for bounds in ([2048, 65, 64, 0, 16], [64, 2048, 0, 65, 63]): + lengths.copy_(torch.tensor(bounds, device="cuda", dtype=torch.int32)) + q.normal_() + graph.replay() + torch.testing.assert_close(actual, compute(False), atol=0, rtol=0) diff --git a/unit_tests/common/qsa/test_flashinfer_topk.py b/unit_tests/common/qsa/test_flashinfer_topk.py new file mode 100644 index 0000000000..5bc432abf1 --- /dev/null +++ b/unit_tests/common/qsa/test_flashinfer_topk.py @@ -0,0 +1,84 @@ +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.qsa.indexer import QSAIndexer +from lightllm.common.qsa.topk import flashinfer_topk_indices, flashinfer_topk_supported + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.fixture(autouse=True) +def require_ragged_topk(): + if not flashinfer_topk_supported(torch.device("cuda", 0), 65542, 512): + pytest.skip("FlashInfer graph-safe TopK on Hopper is required") + + +def test_ragged_graph_uses_visible_lengths_and_chunk_offset(): + rows, columns, k, first = 17, 65542, 512, 5 + scores = torch.randn(rows, columns, device="cuda") + visible = torch.zeros(first + rows + 3, dtype=torch.int32, device="cuda") + lengths = visible[first : first + rows] + limits = torch.tensor([0, 3, 2047, 2048, 2051, 65536, 262144], dtype=torch.int32, device="cuda") + cols = torch.arange(columns, device="cuda") + + def select(): + return flashinfer_topk_indices(scores, k, visible_lengths=visible, first=first) + + select() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + selected = select() + for step in range(12): + lengths.copy_(limits[(torch.arange(rows, device="cuda") + step) % limits.numel()]) + scores.normal_() + if step % 3 == 0: + scores.round_() + scores.masked_fill_(cols[None, :] >= lengths[:, None] // 4, 1e9) + scores.masked_fill_((lengths <= k * 4)[:, None], float("nan")) + graph.replay() + for row, complete in enumerate((lengths // 4).tolist()): + actual = selected[row] + if complete <= k: + expected = torch.arange(k, dtype=torch.int32, device="cuda") + expected[expected >= complete] = -1 + assert torch.equal(actual, expected) + else: + assert bool(((actual >= 0) & (actual < complete)).all()) + assert torch.unique(actual).numel() == k + expected = torch.topk(scores[row, :complete], k).values + values = scores[row, actual.long()].sort(descending=True).values + torch.testing.assert_close(values, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("chunk_size", [3, 64]) +def test_indexer_ragged_topk_preserves_slots_across_graph_transitions(chunk_size): + rows, capacity = 17, 262169 + norm = torch.ones(128, dtype=torch.bfloat16, device="cuda") + original = QSAIndexer(torch.randn(640, 16, dtype=torch.bfloat16, device="cuda"), norm, norm) + optimized = replace(original, use_cuda_topk=True) + cache = torch.zeros(capacity + 3, 128, dtype=torch.bfloat16, device="cuda") + # Binary keys and power-of-two queries give unique scores at the selection boundary. + bits = torch.arange(17, device="cuda") + cache[:, :17] = ((torch.arange(cache.shape[0], device="cuda")[:, None] // 4 >> bits) & 1).to(cache.dtype) + query = torch.zeros(rows, 4, 128, dtype=cache.dtype, device="cuda") + query[:, :, :17] = (2.0 ** bits).to(cache.dtype) + table = torch.arange(capacity, dtype=torch.int32, device="cuda").expand(2, -1).contiguous() + requests = torch.arange(rows, dtype=torch.int32, device="cuda") // 4 % 2 + lengths = torch.full((rows,), 262144, dtype=torch.int32, device="cuda") + data = query, cache, requests, lengths, table + original.select(*data, query_chunk_size=chunk_size) + optimized.select(*data, query_chunk_size=chunk_size) + before, after = torch.cuda.CUDAGraph(), torch.cuda.CUDAGraph() + with torch.cuda.graph(before): + expected = original.select(*data, query_chunk_size=chunk_size) + with torch.cuda.graph(after): + actual = optimized.select(*data, query_chunk_size=chunk_size) + limits = torch.tensor([0, 3, 2047, 2048, 2051, 65536, 262144], dtype=torch.int32, device="cuda") + for step in range(12): + lengths.copy_(limits[(torch.arange(rows, device="cuda") + step) % limits.numel()]) + requests.copy_((torch.arange(rows, device="cuda") + step) % 2) + before.replay() + after.replay() + assert torch.equal(actual.sort(dim=-1).values, expected.sort(dim=-1).values) diff --git a/unit_tests/common/qsa/test_indexer.py b/unit_tests/common/qsa/test_indexer.py new file mode 100644 index 0000000000..d82dee83bf --- /dev/null +++ b/unit_tests/common/qsa/test_indexer.py @@ -0,0 +1,419 @@ +import math + +import pytest +import torch +import torch.nn.functional as F +import triton + +from lightllm.common.qsa.indexer import QSAIndexer +from lightllm.common.qsa.indexer_kernels import _scores + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("minimum", [0, 2]) +def test_grouped_scores_keep_chunk_offsets_short_rows_and_mixed_requests(dtype, minimum): + query = torch.randn(14, 4, 128, device="cuda", dtype=dtype) + cache = torch.randn(204, 128, device="cuda", dtype=dtype) + table = torch.randperm(204, device="cuda").int().view(3, 68) + requests = torch.tensor([0] * 9 + [1, 1, 2, 2, 2], device="cuda", dtype=torch.int32) + lengths = torch.tensor([0] * 5 + [0, 3, 8, 16, 7, 9, 12, 17, 11], device="cuda", dtype=torch.int32) + outputs = [torch.full((9, 17), 42.0, device="cuda") for _ in range(2)] + + def run(): + for group, output in zip((1, 4), outputs): + block = 64 if group == 1 else 128 + _scores[(triton.cdiv(9, group), 1)]( + query, + cache, + requests, + lengths, + table, + output, + 5, + 4, + 128, + 17, + 68, + True, + 4, + 16, + 128, + block, + minimum, + "ieee" if dtype == torch.float32 else "tf32", + GROUP_SIZE=group, + TOKENS=9, + ) + + run() + torch.testing.assert_close(outputs[1], outputs[0], atol=1e-5, rtol=1e-5) + if minimum: + assert torch.all(outputs[1][:2] == 42) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + lengths.zero_() + query.normal_() + graph.replay() + torch.testing.assert_close(outputs[1], outputs[0], atol=1e-5, rtol=1e-5) + + +def make_indexer(dtype=torch.bfloat16, dim=32, heads=4, budget=8, ratio=4): + projection = torch.randn((heads + 1) * dim, 32, device="cuda", dtype=dtype) / math.sqrt(32) + return QSAIndexer( + projection, + torch.randn(dim, device="cuda", dtype=dtype) * 0.1, + torch.randn(dim, device="cuda", dtype=dtype) * 0.1, + heads, + budget, + ratio, + ) + + +def rotary(rows, width, dtype): + angles = torch.randn(rows, width // 2, device="cuda").repeat(1, 2) + return angles.cos().to(dtype), angles.sin().to(dtype) + + +def norm(x, weight, eps): + normalized = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps) + return (normalized * (1 + weight.float())).to(x.dtype) + + +def rotate(x, cos, sin): + width = cos.shape[-1] + first, second = x[..., :width].chunk(2, -1) + return torch.cat([x[..., :width] * cos + torch.cat([-second, first], -1) * sin, x[..., width:]], -1) + + +def reference_cache(indexer, raw, table, lengths, cos, sin): + result = indexer.empty_cache(table.shape[0], table.shape[1]).fill_(float("nan")) + for row, length in enumerate(lengths): + count = length // indexer.compress_ratio + if not count: + continue + slots = table[row, : count * indexer.compress_ratio].view(count, indexer.compress_ratio).long() + pooled = raw[slots].float().mean(1).to(raw.dtype) + normalized = norm(pooled, indexer.key_norm, indexer.eps) + result[row, :count] = rotate(normalized, cos[slots[:, 0]], sin[slots[:, 0]]) + return result + + +def reference_selection(indexer, query, cache, requests, lengths, table): + rows = [] + for row, (request, length) in enumerate(zip(requests.tolist(), lengths.tolist())): + count = length // indexer.compress_ratio + if count: + scores = (query[row].float() @ cache[request, :count].float().T).relu().sum(0) / math.sqrt(indexer.head_dim) + blocks = scores.topk(min(count, indexer.token_budget // indexer.compress_ratio)).indices + positions = ( + blocks[:, None] * indexer.compress_ratio + torch.arange(indexer.compress_ratio, device=query.device) + ).flatten() + else: + positions = torch.empty(0, device=query.device, dtype=torch.int64) + positions = torch.cat([positions, torch.arange(count * indexer.compress_ratio, length, device=query.device)]) + selected = table[request, positions] + rows.append( + F.pad(selected, (0, indexer.token_budget + indexer.compress_ratio - 1 - selected.numel()), value=-1) + ) + return torch.stack(rows) if rows else table.new_empty((0, indexer.token_budget + indexer.compress_ratio - 1)) + + +def assert_sets(actual, expected): + torch.testing.assert_close(actual.sort(-1).values, expected.sort(-1).values, atol=0, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("dim,rotary_dim,heads,ratio", [(32, 16, 4, 4), (65, 32, 3, 3), (128, 64, 4, 4)]) +def test_projection_pooling_and_causal_selection(dtype, dim, rotary_dim, heads, ratio): + torch.manual_seed(337) + indexer = make_indexer(dtype, dim, heads, budget=2 * ratio, ratio=ratio) + hidden = torch.randn(10, 32, device="cuda", dtype=dtype) + qcos, qsin = rotary(10, rotary_dim, dtype) + query, raw_new = indexer.project(hidden, qcos, qsin) + projected = F.linear(hidden, indexer.projection) + expected_q = projected[:, : heads * dim].reshape(10, heads, dim) + expected_q = rotate(norm(expected_q, indexer.query_norm, indexer.eps), qcos[:, None], qsin[:, None]) + tolerance = max(2e-6, torch.finfo(dtype).eps * 2) + torch.testing.assert_close(query, expected_q, atol=tolerance, rtol=tolerance) + torch.testing.assert_close(raw_new, projected[:, heads * dim :], atol=0, rtol=0) + assert not raw_new.is_contiguous() + raw_copy = torch.empty_like(raw_new) + raw_copy.index_copy_(0, torch.arange(raw_new.shape[0], device="cuda"), raw_new) + torch.testing.assert_close(raw_copy, raw_new.contiguous(), atol=0, rtol=0) + raw = torch.randn(100, dim, device="cuda", dtype=dtype) + cos, sin = rotary(100, rotary_dim, dtype) + table = torch.randperm(100, device="cuda")[:87].view(3, 29).int() + lengths = [23, 17, 0] + cache = indexer.empty_cache(3, 29).fill_(float("nan")) + raw_before = raw.clone() + indexer.pool( + raw, + table, + torch.tensor(lengths, device="cuda", dtype=torch.int32), + cos, + sin, + cache, + torch.zeros(3, device="cuda", dtype=torch.int32), + 29 // ratio, + ) + expected_cache = reference_cache(indexer, raw, table, lengths, cos, sin) + torch.testing.assert_close(cache, expected_cache, atol=tolerance, rtol=tolerance, equal_nan=True) + torch.testing.assert_close(raw, raw_before, atol=0, rtol=0) + requests = torch.tensor([0, 0, 1, 0, 1, 1, 0, 1, 1, 0], device="cuda", dtype=torch.int32) + visible = torch.tensor([0, 1, 3, 4, 5, 8, 9, 15, 17, 23], device="cuda", dtype=torch.int32) + actual = indexer.select(query, cache, requests, visible, table, query_chunk_size=3) + + expected = reference_selection(indexer, query, cache, requests, visible, table) + assert_sets(actual, expected) + for row, (request, length) in enumerate(zip(requests.tolist(), visible.tolist())): + allowed = set(table[request, :length].tolist()) + assert set(actual[row][actual[row] >= 0].tolist()) <= allowed + + +def test_incremental_pool_cache_and_cuda_graph_replay(): + torch.manual_seed(337) + indexer = make_indexer() + raw = torch.randn(48, 32, device="cuda", dtype=torch.bfloat16) + cos, sin = rotary(48, 16, raw.dtype) + table = torch.randperm(48, device="cuda")[:40].view(2, 20).int() + cache = indexer.empty_cache(2, 20).fill_(float("nan")) + lengths = torch.tensor([7, 5], device="cuda", dtype=torch.int32) + first = torch.zeros_like(lengths) + indexer.pool(raw, table, lengths, cos, sin, cache, first, 2) + original = cache[:, :1].clone() + hidden = torch.randn(2, 32, device="cuda", dtype=raw.dtype) + qcos, qsin = rotary(2, 16, raw.dtype) + requests = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + first.fill_(1) + lengths.copy_(torch.tensor([13, 10], device="cuda", dtype=torch.int32)) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + query, _ = indexer.project(hidden, qcos, qsin) + indexer.pool(raw, table, lengths, cos, sin, cache, first, 2) + indexer.select(query, cache, requests, lengths, table, query_chunk_size=1) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + query, _ = indexer.project(hidden, qcos, qsin) + indexer.pool(raw, table, lengths, cos, sin, cache, first, 2) + selected = indexer.select(query, cache, requests, lengths, table, query_chunk_size=1) + for counts, starts in [([13, 10], [1, 1]), ([18, 15], [3, 2])]: + hidden.normal_() + lengths.copy_(torch.tensor(counts, device="cuda", dtype=torch.int32)) + first.copy_(torch.tensor(starts, device="cuda", dtype=torch.int32)) + graph.replay() + expected_cache = reference_cache(indexer, raw, table, counts, cos, sin) + tolerance = 2 * torch.finfo(raw.dtype).eps + torch.testing.assert_close(cache, expected_cache, atol=tolerance, rtol=tolerance, equal_nan=True) + assert_sets(selected, reference_selection(indexer, query, cache, requests, lengths, table)) + torch.testing.assert_close(cache[:, :1], original, atol=0, rtol=0) + + +def test_zero_score_ties_tail_only_empty_and_checkpoint_binding(): + indexer = make_indexer() + prefix = "model.language_model.layers.3.self_attn.indexer" + weights = { + f"{prefix}.index_qk_proj.weight": indexer.projection, + f"{prefix}.q_layernorm.weight": indexer.query_norm, + f"{prefix}.k_layernorm.weight": indexer.key_norm, + } + bound = QSAIndexer.from_checkpoint(weights, prefix, token_budget=8) + assert bound.projection is indexer.projection + query = torch.zeros(3, 4, 32, device="cuda", dtype=torch.bfloat16) + table = torch.arange(29, device="cuda", dtype=torch.int32).view(1, -1) + cache = torch.randn(1, 7, 32, device="cuda", dtype=query.dtype) + requests = torch.zeros(3, device="cuda", dtype=torch.int32) + lengths = torch.tensor([0, 3, 21], device="cuda", dtype=torch.int32) + selected = bound.select(query, cache, requests, lengths, table) + assert torch.all(selected[0] == -1) + assert set(selected[1][selected[1] >= 0].tolist()) == {0, 1, 2} + chosen = selected[2][selected[2] >= 0].tolist() + assert len(chosen) == len(set(chosen)) == 9 + assert chosen[-1] == 20 + for offset in [0, 4]: + block = chosen[offset : offset + 4] + assert block == list(range(block[0], block[0] + 4)) and block[0] % 4 == 0 + empty_query, raw = bound.project(query.new_empty((0, 32)), query.new_empty((0, 0)), query.new_empty((0, 0))) + assert raw.shape == (0, 32) + assert bound.select(empty_query, cache, requests[:0], lengths[:0], table).shape == (0, 11) + + short = bound.select(query[:1], cache[:, :0].contiguous(), requests[:1], lengths[1:2], table[:, :3].contiguous()) + assert set(short[0][short[0] >= 0].tolist()) == {0, 1, 2} + + +def test_invalid_layouts(): + indexer = make_indexer() + with pytest.raises(ValueError, match="budget"): + QSAIndexer(indexer.projection, indexer.query_norm, indexer.key_norm, token_budget=3) + with pytest.raises(ValueError, match="rotary"): + indexer.project( + indexer.projection.new_empty((2, 32)), + indexer.projection.new_empty((2, 3)), + indexer.projection.new_empty((2, 3)), + ) + with pytest.raises(ValueError, match="int32"): + indexer.select( + indexer.projection.new_empty((2, 4, 32)), + indexer.empty_cache(1, 4), + torch.zeros(2, device="cuda", dtype=torch.int64), + torch.ones(2, device="cuda", dtype=torch.int32), + torch.zeros(1, 4, device="cuda", dtype=torch.int32), + ) + + +def test_rejected_suffix_repool_overwrites_candidate_blocks(): + torch.manual_seed(337) + indexer = make_indexer() + table = torch.arange(20, device="cuda", dtype=torch.int32).view(1, -1) + raw = torch.randn(20, 32, device="cuda", dtype=torch.bfloat16) + cos, sin = rotary(20, 16, raw.dtype) + cache = indexer.empty_cache(1, 20).fill_(float("nan")) + indexer.pool( + raw, + table, + torch.tensor([13], device="cuda", dtype=torch.int32), + cos, + sin, + cache, + torch.tensor([0], device="cuda", dtype=torch.int32), + 3, + ) + prefix = cache[:, :1].clone() + + raw[7:12].normal_() + lengths = torch.tensor([12], device="cuda", dtype=torch.int32) + indexer.pool(raw, table, lengths, cos, sin, cache, torch.tensor([7 // 4], device="cuda", dtype=torch.int32), 2) + expected = reference_cache(indexer, raw, table, [12], cos, sin) + torch.testing.assert_close(cache, expected, atol=0.016, rtol=0.016, equal_nan=True) + torch.testing.assert_close(cache[:, :1], prefix, atol=0, rtol=0) + query = torch.randn(1, 4, 32, device="cuda", dtype=raw.dtype) + requests = torch.zeros(1, device="cuda", dtype=torch.int32) + assert_sets( + indexer.select(query, cache, requests, lengths, table), + reference_selection(indexer, query, cache, requests, lengths, table), + ) + + +def test_million_token_context_with_bounded_score_chunks(): + torch.manual_seed(337) + context = 1048579 + indexer = make_indexer(dim=128, budget=2048) + cache = torch.randn(1, context // 4, 128, device="cuda", dtype=torch.bfloat16) + + table = torch.arange(context - 1, -1, -1, device="cuda", dtype=torch.int32).view(1, -1) + query = torch.randn(3, 4, 128, device="cuda", dtype=torch.bfloat16) + requests = torch.zeros(3, device="cuda", dtype=torch.int32) + lengths = torch.tensor([4099, 262147, context], device="cuda", dtype=torch.int32) + actual = indexer.select(query, cache, requests, lengths, table, query_chunk_size=1) + expected = reference_selection(indexer, query, cache, requests, lengths, table) + assert_sets(actual, expected) + + +def test_active_rows_pool_into_shared_request_table(): + torch.manual_seed(337) + indexer = make_indexer() + table = torch.randperm(80, device="cuda").view(4, 20).int() + raw = torch.randn(80, 32, device="cuda", dtype=torch.bfloat16) + cos, sin = rotary(80, 16, raw.dtype) + cache = indexer.empty_cache(4, 20).fill_(float("nan")) + rows = torch.tensor([3, -1, 1], device="cuda", dtype=torch.int32) + lengths = torch.tensor([13, 0, 9], device="cuda", dtype=torch.int32) + indexer.pool(raw, table, lengths, cos, sin, cache, torch.zeros_like(lengths), 4, request_rows=rows) + expected = reference_cache(indexer, raw, table, [0, 9, 0, 13], cos, sin) + torch.testing.assert_close(cache, expected, atol=0.016, rtol=0.016, equal_nan=True) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("queries_per_request", [1, 4]) +@pytest.mark.parametrize("query_chunk_size", [3, 64]) +def test_physical_cache_partial_prefix_forks_and_graph_replay(dtype, queries_per_request, query_chunk_size): + torch.manual_seed(740) + indexer = make_indexer(dtype) + raw = torch.randn(32, 32, device="cuda", dtype=dtype) + cos, sin = rotary(32, 16, dtype) + + table = torch.tensor( + [[0, 1, 2, 3, 4, 5, 6, 7, 9, 10, 11, 12], [0, 1, 2, 3, 4, 5, 6, 8, 13, 14, 15, 16]], + device="cuda", + dtype=torch.int32, + ) + cache = indexer.empty_physical_cache(raw.shape[0]).fill_(float("nan")) + lengths = torch.tensor([7, 7], device="cuda", dtype=torch.int32) + first = torch.zeros_like(lengths) + rows = torch.tensor([1, -1, 0], device="cuda", dtype=torch.int32) + active_lengths = torch.tensor([7, 0, 7], device="cuda", dtype=torch.int32) + active_first = torch.zeros_like(active_lengths) + query = torch.randn(2 * queries_per_request, indexer.query_heads, indexer.head_dim, device="cuda", dtype=dtype) + requests = torch.arange(2, device="cuda", dtype=torch.int32).repeat_interleave(queries_per_request) + + def run(): + indexer.pool(raw, table, active_lengths, cos, sin, cache, active_first, 3, request_rows=rows) + return indexer.select( + query, + cache, + requests, + lengths.repeat_interleave(queries_per_request), + table, + query_chunk_size=query_chunk_size, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + selected = run() + + for length in (8, 12, 5, 12): + lengths.fill_(length) + active_lengths.copy_(torch.tensor([length, 0, length], device="cuda", dtype=torch.int32)) + graph.replay() + dense = indexer.empty_cache(2, table.shape[1]).fill_(float("nan")) + indexer.pool(raw, table, lengths, cos, sin, dense, first, 3) + for request in range(2): + count = length // indexer.compress_ratio + end_slots = table[ + request, indexer.compress_ratio - 1 : count * indexer.compress_ratio : indexer.compress_ratio + ] + torch.testing.assert_close(cache[end_slots.long()], dense[request, :count], atol=0, rtol=0) + assert_sets( + selected, indexer.select(query, dense, requests, lengths.repeat_interleave(queries_per_request), table) + ) + assert not torch.equal(cache[7], cache[8]) + + untouched = cache[31].clone() + raw[7].add_(3) + run() + dense = indexer.empty_cache(2, table.shape[1]) + indexer.pool(raw, table, lengths, cos, sin, dense, first, 3) + torch.testing.assert_close(cache[7], dense[0, 1], atol=0, rtol=0) + torch.testing.assert_close(cache[31], untouched, equal_nan=True) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("capacity", [4099, 262159]) +def test_sparse_visibility_in_large_cache_and_dynamic_graph(dtype, capacity): + torch.manual_seed(38413) + indexer = make_indexer(dtype, dim=128, heads=4, budget=8) + query = torch.randn(8, 4, 128, device="cuda", dtype=dtype) + table = torch.stack([torch.randperm(capacity, device="cuda") for _ in range(3)]).int() + requests = torch.tensor([0, 0, 0, 0, 1, 1, 2, 2], device="cuda", dtype=torch.int32) + original = torch.tensor([0, 3, 4, 257, 255, 256, capacity - 1, capacity], device="cuda", dtype=torch.int32) + lengths = original.clone() + cache = torch.randn(3, capacity // 4, 128, device="cuda", dtype=dtype) + cache[0, 257 // 4 :] = float("nan") + cache[1, 256 // 4 :] = float("nan") + actual = indexer.select(query, cache, requests, lengths, table) + assert_sets(actual, reference_selection(indexer, query, cache, requests, lengths, table)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = indexer.select(query, cache, requests, lengths, table) + for pattern in (torch.zeros_like(original), original): + lengths.copy_(pattern) + graph.replay() + assert_sets(output, reference_selection(indexer, query, cache, requests, lengths, table)) + graph.reset() diff --git a/unit_tests/common/qsa/test_topk_compatibility.py b/unit_tests/common/qsa/test_topk_compatibility.py new file mode 100644 index 0000000000..05f2e4f36b --- /dev/null +++ b/unit_tests/common/qsa/test_topk_compatibility.py @@ -0,0 +1,142 @@ +from types import SimpleNamespace + +import pytest +import torch + + +from lightllm.common.qsa import topk +from lightllm.common.qsa.indexer import QSAIndexer + + +@pytest.fixture(autouse=True) +def no_cuda_initialization(): + initially_initialized = torch.cuda.is_initialized() + yield + assert torch.cuda.is_initialized() == initially_initialized + + +@pytest.mark.parametrize("shape,k", [((3, 1027), 512), ((0, 1027), 512), ((2, 11), 0), ((2, 11), 11), ((2, 3, 11), 4)]) +def test_cpu_fallback_keeps_actual_torch_indices(shape, k, monkeypatch): + def unexpected_query(device): + raise AssertionError("CPU fallback queried CUDA") + + monkeypatch.setattr(topk, "_supports_device", unexpected_query) + scores = torch.arange(torch.tensor(shape).prod().item()).reshape(shape).float() + assert torch.equal(topk.cuda_topk_indices(scores, k), torch.topk(scores, k, dim=-1, sorted=False).indices) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64, torch.int64]) +def test_noncontiguous_cpu_fallback(dtype): + scores = torch.arange(22).reshape(11, 2).T.to(dtype) + assert not scores.is_contiguous() + assert torch.equal(topk.cuda_topk_indices(scores, 4), torch.topk(scores, 4, sorted=False).indices) + + +def test_ragged_topk_cpu_and_other_budgets_do_not_query_cuda(monkeypatch): + def unexpected_query(device): + pytest.fail("Unqualified ragged TopK must not initialize CUDA") + + monkeypatch.setattr(torch.cuda, "get_device_capability", unexpected_query) + topk.flashinfer_topk_supported.cache_clear() + assert not topk.flashinfer_topk_supported(torch.device("cpu"), 65542, 512) + assert not topk.flashinfer_topk_supported(torch.device("cuda", 0), 65542, 256) + topk.flashinfer_topk_supported.cache_clear() + + +@pytest.mark.parametrize("k", [-1, 12]) +def test_invalid_budget_keeps_torch_exception(k): + with pytest.raises(RuntimeError): + topk.cuda_topk_indices(torch.ones((2, 11)), k) + + +@pytest.mark.parametrize( + "change", + [ + dict(runtime=False), + dict(dtype=torch.bfloat16), + dict(shape=(3, 1024)), + dict(shape=(3, 8, 1027)), + dict(k=256), + dict(contiguous=False), + dict(device_supported=False), + ], +) +def test_unqualified_cuda_metadata_delegates_without_launch(change, monkeypatch): + config = dict(runtime=True, dtype=torch.float32, shape=(3, 1027), k=512, contiguous=True, device_supported=True) + config.update(change) + scores = SimpleNamespace( + is_cuda=True, + dtype=config["dtype"], + ndim=len(config["shape"]), + shape=config["shape"], + device=torch.device("cuda", 0), + is_contiguous=lambda: config["contiguous"], + ) + sentinel = object() + calls = [] + monkeypatch.setattr(topk, "_CUDA_TOPK_RUNTIME", config["runtime"]) + monkeypatch.setattr(topk, "_supports_device", lambda _: config["device_supported"]) + + def original(value, k, *, dim, sorted): + calls.append((value, k, dim, sorted)) + return SimpleNamespace(indices=sentinel) + + monkeypatch.setattr(topk.torch, "topk", original) + assert topk.cuda_topk_indices(scores, config["k"]) is sentinel + assert calls == [(scores, config["k"], -1, False)] + + +def test_shared_default_and_select_api_remain_compatible(): + import inspect + + assert QSAIndexer.__dataclass_fields__["use_cuda_topk"].default is False + signature = inspect.signature(QSAIndexer.select) + assert list(signature.parameters) == [ + "self", + "query", + "cache", + "requests", + "visible_lengths", + "indices", + "query_chunk_size", + ] + assert signature.parameters["query_chunk_size"].default == 64 + + +@pytest.mark.parametrize( + "change", + [ + dict(dtype=torch.int64), + dict(device=torch.device("cpu")), + dict(ndim=2), + dict(contiguous=False), + dict(first=-1), + dict(first=2), + dict(ratio=0), + ], +) +def test_invalid_visible_metadata_fails_before_allocation(change, monkeypatch): + device = torch.device("cuda", 0) + config = dict(dtype=torch.int32, device=device, ndim=1, contiguous=True, first=0, ratio=4) + config.update(change) + scores = SimpleNamespace( + is_cuda=True, dtype=torch.float32, ndim=2, shape=(3, 1027), device=device, is_contiguous=lambda: True + ) + visible = SimpleNamespace( + dtype=config["dtype"], + device=config["device"], + ndim=config["ndim"], + is_contiguous=lambda: config["contiguous"], + numel=lambda: 3, + ) + monkeypatch.setattr(topk, "_CUDA_TOPK_RUNTIME", True) + monkeypatch.setattr(topk, "_supports_device", lambda _: True) + + def unexpected_allocation(*args, **kwargs): + raise AssertionError("Invalid visible metadata reached GPU allocation") + + monkeypatch.setattr(topk.torch, "empty", unexpected_allocation) + with pytest.raises(ValueError, match="visible lengths"): + topk.cuda_topk_indices( + scores, 512, visible_lengths=visible, first=config["first"], compress_ratio=config["ratio"] + ) diff --git a/unit_tests/common/qsa/test_topk_cuda.py b/unit_tests/common/qsa/test_topk_cuda.py new file mode 100644 index 0000000000..e36249ce3a --- /dev/null +++ b/unit_tests/common/qsa/test_topk_cuda.py @@ -0,0 +1,62 @@ +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.qsa.indexer import QSAIndexer +from lightllm.common.qsa.topk import cuda_topk_supported + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def components(rows): + if not cuda_topk_supported(torch.device("cuda", 0), 1027, 512): + pytest.skip("The optional selector requires its qualified H100 runtime") + weights = torch.randn((640, 16), device="cuda", dtype=torch.bfloat16) + norm = torch.ones(128, device="cuda", dtype=torch.bfloat16) + original = QSAIndexer(weights, norm, norm) + optimized = replace(original, use_cuda_topk=True) + query = torch.randn((rows, 4, 128), device="cuda", dtype=torch.bfloat16) + cache = torch.randn((16384, 128), device="cuda", dtype=torch.bfloat16) + requests = torch.arange(rows, device="cuda", dtype=torch.int32) % 9 + lengths = torch.zeros(rows, device="cuda", dtype=torch.int32) + table = torch.randint(0, 16384, (9, 4110), device="cuda", dtype=torch.int32) + return original, optimized, (query, cache, requests, lengths, table) + + +@pytest.mark.parametrize("rows", [3, 24]) +@pytest.mark.parametrize("visible", [0, 3, 2047, 2048, 2051, 4096]) +def test_optional_selector_preserves_unsorted_slot_order(rows, visible): + original, optimized, data = components(rows) + data[3].fill_(visible) + assert torch.equal(original.select(*data), optimized.select(*data)) + + +@pytest.mark.parametrize("rows", [3, 24]) +def test_short_long_graph_transitions_ignore_poisoned_short_scores(rows, monkeypatch): + original, optimized, data = components(rows) + allocate = torch.empty + + def poisoned_empty(*args, **kwargs): + value = allocate(*args, **kwargs) + if value.is_cuda and value.dtype == torch.float32 and value.ndim == 2 and value.shape[1] == 1027: + value.fill_(float("nan")) + return value + + monkeypatch.setattr(torch, "empty", poisoned_empty) + data[3].fill_(4096) + original.select(*data) + optimized.select(*data) + torch.cuda.synchronize() + before, after = torch.cuda.CUDAGraph(), torch.cuda.CUDAGraph() + with torch.cuda.graph(before): + expected = original.select(*data) + with torch.cuda.graph(after): + actual = optimized.select(*data) + limits = torch.tensor([0, 3, 2047, 2048, 2051, 4096], device="cuda", dtype=torch.int32) + for step in range(24): + data[0].normal_() + data[3].copy_(limits[(torch.arange(rows, device="cuda") + step) % limits.numel()]) + before.replay() + after.replay() + assert torch.equal(actual, expected) diff --git a/unit_tests/common/quantization/test_aligned_projection.py b/unit_tests/common/quantization/test_aligned_projection.py new file mode 100644 index 0000000000..e0c955c58d --- /dev/null +++ b/unit_tests/common/quantization/test_aligned_projection.py @@ -0,0 +1,97 @@ +import pytest +import torch + +from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager +from lightllm.common.quantization.aligned_projection import OutputAlignedProjection + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.fixture(autouse=True) +def current_device(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_unaligned_projection_preserves_loading_layout_outputs_and_graph_updates(dtype): + torch.manual_seed(3826) + dimensions = [256, 256, 768, 768, 6, 6] + method = OutputAlignedProjection() + pack, parts = method.create_weight(dimensions, 2560, dtype, 0) + weights = [torch.randn(width, 2560, device="cuda", dtype=dtype) * 0.02 for width in dimensions] + for source, part in zip(weights, parts): + method.load_weight(source, part) + assert pack.weight.shape == (2060, 2560) + assert torch.equal(pack.weight, torch.cat(weights)) + assert method._padded_weight.shape == (2064, 2560) + assert not torch.count_nonzero(method._padded_weight[2060:]) + + hidden = torch.randn(24, 2560, device="cuda", dtype=dtype) + + def project(value, out=None): + return method.apply(value, pack, out=out, use_custom_tensor_mananger=False) + + singles = torch.cat([project(row[None]) for row in hidden]) + for count in (0, 1, 2, 3, 8, 9, 24): + actual = project(hidden[:count]) + assert actual.shape == (count, 2060) and actual.is_contiguous() + torch.testing.assert_close(actual, singles[:count], atol=0, rtol=0) + out = torch.empty(count, 4120, device="cuda", dtype=dtype)[:, ::2] + assert project(hidden[:count], out) is out + torch.testing.assert_close(out, actual, atol=0, rtol=0) + + inputs = hidden[:3].clone() + for _ in range(3): + project(inputs) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = project(inputs) + inputs.add_(0.125) + graph.replay() + torch.testing.assert_close(captured, project(inputs), atol=0, rtol=0) + method.load_weight(weights[-1] + 0.25, parts[-1]) + graph.replay() + torch.testing.assert_close(captured, project(inputs), atol=0, rtol=0) + assert not torch.count_nonzero(method._padded_weight[2060:]) + graph.reset() + + +@pytest.mark.parametrize("dtype,width", [(torch.bfloat16, 2064), (torch.float32, 2060)]) +def test_aligned_and_fp32_paths_keep_native_projection(dtype, width): + method = OutputAlignedProjection() + pack, parts = method.create_weight([width], 128, dtype, 0) + source = torch.randn(width, 128, device="cuda", dtype=dtype) + method.load_weight(source, parts[0]) + assert method._padded_weight is None + hidden = torch.randn(3, 128, device="cuda", dtype=dtype) + result = method.apply(hidden, pack, use_custom_tensor_mananger=False) + torch.testing.assert_close(result, hidden @ source.T, atol=0, rtol=0) + + +def test_cached_outputs_and_graph_replay_keep_live_buffers_independent(): + method = OutputAlignedProjection() + pack, parts = method.create_weight([2060], 2560, torch.bfloat16, 0) + source = torch.randn(2060, 2560, device="cuda", dtype=torch.bfloat16) * 0.02 + method.load_weight(source, parts[0]) + inputs = torch.randn(3, 2560, device="cuda", dtype=torch.bfloat16) + expected = method.apply(inputs, pack, use_custom_tensor_mananger=False) + graph = torch.cuda.CUDAGraph() + g_cache_manager.cache_env_in() + try: + live = method.apply(inputs, pack) + for _ in range(4): + other = method.apply(inputs + 0.25, pack) + assert other.data_ptr() != live.data_ptr() + torch.testing.assert_close(live, expected, atol=0, rtol=0) + del other + with torch.cuda.graph(graph): + captured = method.apply(inputs, pack) + inputs.add_(0.125) + graph.replay() + torch.testing.assert_close( + captured, method.apply(inputs, pack, use_custom_tensor_mananger=False), atol=0, rtol=0 + ) + torch.testing.assert_close(live, expected, atol=0, rtol=0) + finally: + graph.reset() + g_cache_manager.cache_env_out() diff --git a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py new file mode 100644 index 0000000000..3f937a11ee --- /dev/null +++ b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py @@ -0,0 +1,144 @@ +import multiprocessing +import traceback +from multiprocessing.reduction import ForkingPickler + +import pytest +import torch + +from lightllm.common.attention_state.bounded_att import GDNState, KDAState +from lightllm.server.router.model_infer.mode_backend.pd.p2p_fix import reduce_tensor +from unit_tests.common.test_linear_state import make_state + + +class _SerializationFailure: + def __reduce_ex__(self, protocol): + raise RuntimeError("injected serialization failure") + + +def _assert_received_manager_roundtrip(payload, state_name, mode, expected): + resources = ForkingPickler.loads(payload) + state = getattr(resources, state_name) + assert state.big_pages_pool is state.small_pages_pool is None + assert resources.full_att_allocator is None + layout = resources.get_pd_operator().get_pd_bounded_att_page_layouts()[state_name] + block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + state.write_to_pd_page(1, block, 0, 8, layout) + state.read_from_pd_page(block, 2, 0, 8, layout) + if hasattr(state, "get_accepted_state"): + actual = state.get_accepted_state(2) + else: + width = state.config.conv_width + actual = (state.conv_state[:, 2, :, :width], state.ssm_state[:, 2 * (state.mtp_step + 1)]) + for tensor, values in zip(actual, expected): + torch.testing.assert_close(tensor, torch.tensor(values, dtype=tensor.dtype, device="cuda"), rtol=0, atol=0) + assert state.mtp_state_index[2] == 0 + assert state.mtp_state_index[1] == 1 + if mode == "replay": + assert state.cursors[2] == 0 + assert state.cursors[1] == 1 + torch.cuda.synchronize() + + +def _receive_pd_managers(connection): + torch.cuda.set_device(0) + while True: + message = connection.recv() + if message is None: + return + try: + _assert_received_manager_roundtrip(*message) + connection.send(None) + except Exception: + connection.send(traceback.format_exc()) + + +@pytest.fixture(scope="module") +def pd_receiver(): + context = multiprocessing.get_context("spawn") + sender, receiver = context.Pipe() + process = context.Process(target=_receive_pd_managers, args=(receiver,)) + process.start() + receiver.close() + try: + yield sender, process + finally: + if process.is_alive(): + sender.send(None) + process.join(timeout=30) + if process.is_alive(): + process.terminate() + process.join(timeout=5) + sender.close() + assert process.exitcode == 0 + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mode", ["native", "replay"]) +def test_pd_export_copy_keeps_worker_pools_and_runtime_identity(cls, mode): + state, resources, _ = make_state(cls=cls, mode=mode, device="meta") + state.big_pages_pool, state.small_pages_pool = object(), object() + before = vars(state).copy() + exported = resources.get_pd_operator().create_ipc_copy() + copy = getattr(exported, state.state_name) + assert copy.big_pages_pool is copy.small_pages_pool is None + assert copy.conv_state is state.conv_state + assert copy.ssm_state is state.ssm_state + if mode == "replay": + assert copy.cursors is state.cursors + assert copy.keys is state.keys + assert copy.mtp_state_index is state.mtp_state_index + assert all(vars(state)[name] is value for name, value in before.items()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA IPC and pinned memory required") +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mode", ["native", "replay"]) +@pytest.mark.parametrize("fail_serialization", [False, True]) +def test_shared_attention_state_roundtrip_preserves_worker_resources( + runtime_start_args, monkeypatch, pd_receiver, cls, mode, fail_serialization +): + state, resources, _ = make_state(cls=cls, mode=mode) + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) + state.conv_state.copy_(torch.arange(5, dtype=torch.bfloat16, device="cuda")) + for slot in range(state.ssm_state.shape[1]): + state.ssm_state[:, slot].fill_(slot + 1) + state.mtp_state_index[1], state.mtp_state_index[2] = 1, 2 + if mode == "replay": + state.raw_keys.zero_() + state.raw_values.zero_() + state.gates.fill_(-0.5) + state.betas.zero_() + state.cursors[1], state.cursors[2] = 1, 3 + if hasattr(state, "get_accepted_state"): + expected = [tensor.cpu().tolist() for tensor in state.get_accepted_state(1)] + else: + width = state.config.conv_width + expected = [ + state.conv_state[:, 1, :, :width].cpu().tolist(), + state.ssm_state[:, state.mtp_step + 1].cpu().tolist(), + ] + before = vars(state).copy() + source_slot = 1 if mode == "replay" else state.mtp_step + 1 + source_conv, source_ssm = state.conv_state[:, 1].clone(), state.ssm_state[:, source_slot].clone() + snapshots = [] + for pool in (state.big_pages_pool, state.small_pages_pool): + for tensor in (pool.conv_state, pool.ssm_state): + tensor.fill_(3) + snapshots.append((tensor, tensor.data_ptr(), tensor.clone())) + monkeypatch.setitem(ForkingPickler._extra_reducers, torch.Tensor, reduce_tensor) + exported = resources.get_pd_operator().create_ipc_copy() + if fail_serialization: + with pytest.raises(RuntimeError, match="injected serialization failure"): + ForkingPickler.dumps((_SerializationFailure(), exported)) + else: + sender, receiver = pd_receiver + sender.send((ForkingPickler.dumps(exported).tobytes(), state.state_name, mode, expected)) + assert sender.poll(30), f"PD receiver did not respond; exit code: {receiver.exitcode}" + error = sender.recv() + assert error is None, error + assert all(vars(state)[name] is value for name, value in before.items()) + torch.testing.assert_close(state.conv_state[:, 1], source_conv, rtol=0, atol=0) + torch.testing.assert_close(state.ssm_state[:, source_slot], source_ssm, rtol=0, atol=0) + for tensor, pointer, contents in snapshots: + assert tensor.data_ptr() == pointer + torch.testing.assert_close(tensor, contents, rtol=0, atol=0) diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py new file mode 100644 index 0000000000..a3a11d0b7b --- /dev/null +++ b/unit_tests/common/test_linear_state.py @@ -0,0 +1,268 @@ +from dataclasses import replace +from importlib import import_module + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState, KDAState + + +def make_state(cls=GDNState, mode="native", step=2, dtype=torch.float32, device="cuda", config=None): + config = config or ConvSSMStateConfig( + logical_layer_ids=(0, 1), + global_k_heads=1, + global_v_heads=1 if cls is KDAState else 2, + key_dim=32, + value_dim=32, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=dtype, + ) + if mode == "replay": + 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 + + cls = ReplayKDAState if cls is KDAState else ReplayGDNState + state = cls(config, 4, mtp_step=step, device=device, ssm_replay_config={"cache_len": 4}) + else: + state = cls(config, 4, mtp_step=step, device=device) + resources = AttentionStateManager(request_capacity=4) + resources.register_state("main", state) + resources.finish_state_registration() + layout = resources.get_pd_operator().get_pd_bounded_att_page_layouts()[state.state_name] + return state, resources, layout + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mode,step", [("native", 0), ("native", 2), ("replay", 0), ("replay", 2)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( + runtime_start_args, cls, mode, step, dtype +): + state, _, layout = make_state(cls, mode, step, dtype) + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) + src, dst, offset = 1, 2, min(1, step) + state.conv_state.copy_(torch.arange(3 + step, device="cuda").to(torch.bfloat16)) + for slot in range(state.ssm_state.shape[1]): + state.ssm_state[:, slot].fill_(slot + 1) + if state.mtp_state_index is not None: + state.mtp_state_index[src] = offset + if mode == "native": + expected_conv = state.conv_state[:, src, ..., :3].clone() + expected_ssm = state.ssm_state[:, src * (state.mtp_step + 1)].clone() + else: + expected_conv = state.conv_state[:, src, ..., offset : offset + 3].clone() + expected_ssm = state.ssm_state[:, src].clone() + if mode == "replay": + state.raw_keys.zero_() + state.raw_values.zero_() + state.gates.fill_(-0.5) + state.betas.zero_() + state.cursors[src] = 1 + expected_ssm = (expected_ssm.float() * torch.exp(torch.tensor(-0.5))).to(dtype) + before = state.ssm_state.clone() + block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + state.write_to_pd_page(src, block, 0, 8, layout) + for kind in ("small", "big"): + state.save_bounded_att_runtime_to_pages([src], [8], kind, [0]) + torch.cuda.synchronize() + for pool in (state.small_pages_pool, state.big_pages_pool): + torch.testing.assert_close(pool.conv_state[0], expected_conv.cpu(), rtol=0, atol=0) + torch.testing.assert_close(pool.ssm_state[0], expected_ssm.cpu(), rtol=2e-5, atol=2e-5) + assert torch.equal(state.ssm_state, before) + if mode == "replay": + assert state.cursors[src].item() == 1 + + for restore in ( + lambda: state.restore_bounded_att_runtime_from_pages([dst], [8], "small", [0]), + lambda: state.restore_bounded_att_runtime_from_pages([dst], [8], "big", [0]), + lambda: state.read_from_pd_page(block, dst, 0, 8, layout), + ): + if state.mtp_state_index is not None: + state.mtp_state_index[dst] = step + if mode == "replay": + state.cursors[dst] = 3 + restore() + if mode == "replay": + actual = state.get_accepted_state(dst) + actual_conv, actual_ssm = actual.conv, actual.ssm + else: + actual_conv = state.conv_state[:, dst, ..., :3] + actual_ssm = state.ssm_state[:, dst * (state.mtp_step + 1)] + torch.testing.assert_close(actual_conv, expected_conv, rtol=0, atol=0) + torch.testing.assert_close(actual_ssm, expected_ssm, rtol=2e-5, atol=2e-5) + if state.mtp_state_index is not None: + assert state.mtp_state_index[dst].item() == 0 + if mode == "replay": + assert state.cursors[dst].item() == 0 + state.reset_bounded_att_state([dst]) + assert not state.conv_state[:, dst].any() + if mode == "replay": + assert not state.ssm_state[:, dst].any() + else: + start = dst * (state.mtp_step + 1) + assert not state.ssm_state[:, start : start + state.mtp_step + 1].any() + + +def copy_on_cpu(*, copies): + for source, destination in copies: + destination.copy_(source) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mode", ["native", "replay"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("destination", ["cpu", "cuda"]) +def test_checkpoint_batch_export(runtime_start_args, monkeypatch, cls, mode, dtype, destination): + state, _, _ = make_state(cls, mode, 2, dtype) + state.initialize_page_pools(3, device=destination) + state.conv_state.copy_(torch.arange(5, device="cuda").to(torch.bfloat16)) + for slot in range(state.ssm_state.shape[1]): + state.ssm_state[:, slot].fill_(slot + 1) + state.mtp_state_index.copy_(torch.tensor([0, 2, 1, 0], device="cuda")) + if mode == "replay": + state.raw_keys.zero_() + state.raw_values.zero_() + state.gates.fill_(-0.5) + state.betas.zero_() + state.cursors.copy_(torch.tensor([1, 2, 3, 0], device="cuda")) + requests = [2, 0, 1] + if mode == "replay": + expected = [(item.conv.clone(), item.ssm.clone()) for item in map(state.get_accepted_state, requests)] + else: + expected = [ + ( + state.conv_state[:, req, :, : state.config.conv_width].clone(), + state.ssm_state[:, req * (state.mtp_step + 1)].clone(), + ) + for req in requests + ] + before = state.ssm_state.clone() + history = getattr(state, "runtime", state) + cursors = history.cursors.clone() if mode == "replay" else None + + def no_snapshot(*args): + raise AssertionError("batch export must write directly to the destination") + + if mode == "replay": + monkeypatch.setattr(state, "get_accepted_state", no_snapshot) + if history is not state: + monkeypatch.setattr(history, "get_accepted_state", no_snapshot) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + state.save_bounded_att_runtime_to_pages(requests, [8] * 3, "big", [1, 2, 0]) + stream.synchronize() + pool = state.big_pages_pool + for slot, (conv, ssm) in zip([1, 2, 0], expected): + torch.testing.assert_close(pool.conv_state[slot].cpu(), conv.cpu(), atol=0, rtol=0) + torch.testing.assert_close(pool.ssm_state[slot].cpu(), ssm.cpu(), atol=0, rtol=0) + assert torch.equal(state.ssm_state, before) + if cursors is not None: + assert torch.equal(state.cursors, cursors) + target, _, _ = make_state(cls, mode, 2, dtype) + if target.mtp_state_index is not None: + target.mtp_state_index.fill_(2) + target.copy_runtime_state_from_dp(requests, [0, 1, 2], state) + for req, (conv, ssm) in enumerate(expected): + if mode == "replay": + actual = target.get_accepted_state(req) + actual_conv, actual_ssm = actual.conv, actual.ssm + else: + actual_conv = target.conv_state[:, req, :, : target.config.conv_width] + actual_ssm = target.ssm_state[:, req * (target.mtp_step + 1)] + torch.testing.assert_close(actual_conv, conv, atol=0, rtol=0) + torch.testing.assert_close(actual_ssm, ssm, atol=0, rtol=0) + assert torch.equal(state.ssm_state, before) + if cursors is not None: + assert torch.equal(state.cursors, cursors) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_native_checkpoint_scratch_is_bounded(runtime_start_args): + config = ConvSSMStateConfig( + logical_layer_ids=(0, 1, 2, 3), + global_k_heads=1, + global_v_heads=2, + key_dim=128, + value_dim=128, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ) + state = GDNState(config, 65, mtp_step=2, device="cuda") + state.initialize_page_pools(64) + requests = list(range(64)) + state.save_bounded_att_runtime_to_pages(requests, [128] * 64, "big", requests) + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + state.save_bounded_att_runtime_to_pages(requests, [128] * 64, "big", requests) + torch.cuda.synchronize() + assert torch.cuda.max_memory_allocated() - before < 1024 * 1024 + + +def test_pd_block_roundtrip_preserves_replicated_heads_and_accepted_mtp_state(monkeypatch): + # Exercise layout and accepted-state selection independently of the CUDA copy kernel. + monkeypatch.setattr( + import_module("lightllm.common.attention_state.bounded_att.gdn.gdn_state"), + "copy_bounded_att_tensors", + copy_on_cpu, + ) + config = ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=2, + key_dim=2, + value_dim=3, + conv_width=3, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + tp_world_size=4, + ) + states, layouts, snapshots = [], [], [] + src, dst = 1, 2 + for rank in range(4): + state, _, layout = make_state(device="cpu", config=replace(config, tp_rank=rank)) + state.conv_state.fill_(-1) + state.ssm_state.fill_(-2) + state.mtp_state_index.fill_(1) + state.conv_state[:, src, :2, :3].fill_(10 + rank) + state.conv_state[:, src, 2:4, :3].fill_(20 + rank) + state.conv_state[:, src, 4:, :3].fill_(30 + rank) + state.ssm_state[:, src * (state.mtp_step + 1)].fill_(40 + rank) + states.append(state) + layouts.append(layout) + snapshots.append((state.conv_state.clone(), state.ssm_state.clone())) + size = layouts[0].byte_end + block = torch.full((1, size + 16), 0xA5, dtype=torch.uint8) + for state, layout in zip(states, layouts): + state.write_to_pd_page(src, block, 0, 8, layout) + expected_conv = torch.cat( + [torch.full((1, channels, 3), value) for channels, value in ((2, 10), (2, 20), (3, 30), (3, 32))], + dim=1, + ).float() + expected_ssm = torch.stack((torch.full((1, 2, 3), 40), torch.full((1, 2, 3), 42)), dim=1).float() + conv_bytes = expected_conv.numel() * 4 + torch.testing.assert_close(block[0, :conv_bytes].view(torch.float32).view_as(expected_conv), expected_conv) + torch.testing.assert_close(block[0, conv_bytes:size].view(torch.float32).view_as(expected_ssm), expected_ssm) + assert torch.all(block[0, size:] == 0xA5) + exported = block.clone() + other_requests = torch.tensor([0, 1, 3]) + for rank, (state, layout, (conv_before, ssm_before)) in enumerate(zip(states, layouts, snapshots)): + torch.testing.assert_close(state.conv_state, conv_before) + torch.testing.assert_close(state.ssm_state, ssm_before) + state.read_from_pd_page(block, dst, 0, 8, layout) + expected_local_conv = torch.cat( + (expected_conv[:, :4], expected_conv[:, 4 + rank // 2 * 3 : 7 + rank // 2 * 3]), dim=1 + ) + torch.testing.assert_close(state.conv_state[:, dst, ..., :3], expected_local_conv) + torch.testing.assert_close(state.ssm_state[:, dst * 3], expected_ssm[:, rank // 2 : rank // 2 + 1]) + assert state.mtp_state_index.tolist() == [1, 1, 0, 1] + torch.testing.assert_close(state.conv_state[:, other_requests], conv_before[:, other_requests]) + other_slots = (other_requests[:, None] * 3 + torch.arange(3)).flatten() + torch.testing.assert_close(state.ssm_state[:, other_slots], ssm_before[:, other_slots]) + assert torch.equal(block, exported) diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py new file mode 100644 index 0000000000..f720e01f09 --- /dev/null +++ b/unit_tests/common/test_linear_state_contract.py @@ -0,0 +1,201 @@ +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.ssm_autotune import resolve_run_config +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 + + +@pytest.fixture +def linear_config(): + return ConvSSMStateConfig( + logical_layer_ids=(0, 1, 3, 4, 6, 7), + global_k_heads=1, + global_v_heads=2, + key_dim=32, + value_dim=32, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ) + + +@pytest.mark.parametrize("layer,row", [(0, 0), (1, 1), (3, 2), (4, 3), (6, 4), (7, 5)]) +def test_layer_buffers_follow_packed_linear_layers(linear_config, layer, row): + owner = GDNState(linear_config, request_capacity=1, device="cpu") + for i in range(6): + owner.conv_state[i].fill_(i) + owner.ssm_state[i].fill_(i + 10) + actual_conv, actual_ssm = owner.get_layer_view(layer) + assert (actual_conv == row).all() + assert (actual_ssm == row + 10).all() + + +@pytest.mark.parametrize("layer", [-1, 2, 5, 8]) +def test_linear_mapping_rejects_full_attention_and_out_of_range_layers(linear_config, layer): + owner = GDNState(linear_config, request_capacity=1, device="cpu") + with pytest.raises(KeyError): + owner.get_layer_view(layer) + + +def test_model_configuration_retains_nondefault_kda_gate_bound(): + config = Glm5NextResourceConfig.from_model_config( + { + "kv_lora_rank": 512, + "num_hidden_layers": 4, + "layer_types": ["linear_attention"] * 3 + ["deepseek_sparse_attention"], + "index_kpool": 4, + "index_head_dim": 128, + "linear_attn_config": { + "num_heads": 8, + "head_dim": 128, + "short_conv_kernel_size": 4, + "kda_layers": [0, 1, 2], + "gate_lower_bound": -2.5, + }, + }, + StartArgs(tp=4, data_type="bfloat16"), + ) + assert config.gate_lower_bound == -2.5 + + +def test_kda_backend_receives_gate_bound_at_construction(linear_config): + owner = ReplayKDAState(linear_config, 3, mtp_step=2, gate_lower_bound=-2.5, device="cpu") + assert owner.gate_lower_bound == -2.5 + + +def test_replay_backend_publishes_detached_run_config(): + ssm_backend = object.__new__(ReplayGDNState) + tuning = {"BV": 32, "num_warps": 2} + ssm_backend.run_config = tuning + ssm_backend._config_is_fixed = True + ssm_backend.compute_mode = "inline" + config = resolve_run_config(ssm_backend, "replay", *([None] * 7), None) + tuning["BV"] = 128 + assert config == {"BV": 32, "num_warps": 2, "num_stages": 3} + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mtp_step,expected", [(0, "inline"), (2, "precompute")]) +def test_replay_defaults_to_precompute_for_speculative_decode(linear_config, cls, mtp_step, expected): + replay_cls = ReplayKDAState if cls is KDAState else ReplayGDNState + owner = replay_cls(linear_config, request_capacity=2, mtp_step=mtp_step, device="cpu") + history = getattr(owner, "runtime", owner) + assert history.compute_mode == expected + assert history.verify_width == mtp_step + 1 + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +def test_mtp0_decode_keeps_explicit_precompute(linear_config, cls): + replay_cls = ReplayKDAState if cls is KDAState else ReplayGDNState + owner = replay_cls( + linear_config, + request_capacity=2, + mtp_step=0, + device="cpu", + ssm_replay_config={"cache_len": 8, "compute_mode": "precompute"}, + ) + assert owner.verify_width == 1 + assert owner.compute_mode == "precompute" + assert owner.capacity == 8 + + +class _CaptureKernel: + def __init__(self): + self.args = None + + def __getitem__(self, grid): + def launch(*args, **kwargs): + self.args = args + + return launch + + +@pytest.mark.parametrize( + "module_name,cls,bv,precompute", + [ + ("replay_gdn_decode", ReplayGDNState, 32, True), + ("replay_gdn_decode", ReplayGDNState, 8, False), + ("replay_kda_decode", ReplayKDAState, 32, True), + ("replay_kda_decode", ReplayKDAState, 8, False), + ], +) +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}") + captured = _CaptureKernel() + monkeypatch.setattr(decode, "_decode_replay", captured) + owner = cls( + linear_config, + request_capacity=2, + mtp_step=0, + device="cpu", + ssm_replay_config={"cache_len": 8, "compute_mode": "precompute"}, + ) + heads = owner.num_key_heads + q = torch.zeros(1, 1, heads, linear_config.key_dim) + v = torch.zeros(1, 1, owner.ssm_state.shape[-3], linear_config.value_dim) + gate = v.shape[-3] * linear_config.key_dim if cls is ReplayKDAState else v.shape[-3] + decode.replay( + owner, + 0, + q, + q, + v, + torch.zeros(1, gate), + torch.zeros(1, v.shape[-3]), + torch.zeros(v.shape[-3]), + torch.zeros(gate), + torch.zeros(1, dtype=torch.int32), + None, + run_config={"BV": bv, "num_warps": 1, "num_stages": 1}, + ) + assert captured.args[33] is precompute + assert captured.args[35] is (cls is ReplayKDAState) + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +def test_native_commit_selects_mtp_slot(cls): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + owner, _, _ = make_state(cls, step=2) + owner.commit_bounded_att_state( + torch.tensor([1], device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda") + ) + assert owner.mtp_state_index.tolist() == [0, 2, 0, 0] + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +def test_replay_fused_acceptance_keeps_cursors_inside_mode(monkeypatch, cls): + import sys + + owner, _, _ = make_state(cls, mode="replay", step=3, device="cpu") + history = getattr(owner, "runtime", owner) + calls = [] + monkeypatch.setattr(sys.modules[type(history).__module__], "commit_ssm_prefix", lambda *args: calls.append(args)) + monkeypatch.setattr(history, "accept_updates", lambda *a: pytest.fail("unfused acceptance")) + owner.commit_bounded_att_state(torch.tensor([2, 0, 1]), torch.tensor([1, 2, 3])) + assert len(calls) == 1 + assert calls[0][3] is history.cursors + assert calls[0][0] is owner.mtp_state_index + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_acceptance_keeps_request_order_with_strided_inputs(device): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + owner, _, _ = make_state(mode="replay", step=3, device=device) + owner.mtp_state_index.fill_(2) + history = getattr(owner, "runtime", owner) + history.cursors.fill_(8) + reqs = torch.tensor([2, -99, 0, -99, 3, -99], device=device)[::2] + lengths = torch.tensor([1, -99, 3, -99, 2, -99], dtype=torch.int32, device=device)[::2] + owner.commit_accepted(owner.mtp_state_index, reqs, lengths) + assert owner.mtp_state_index.tolist() == [2, 2, 0, 1] + assert history.cursors.tolist() == [11, 8, 9, 8] diff --git a/unit_tests/conftest.py b/unit_tests/conftest.py index 1da9c1e286..622c8087ad 100644 --- a/unit_tests/conftest.py +++ b/unit_tests/conftest.py @@ -22,6 +22,19 @@ def single_rank_environment(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv(key, value) +@pytest.fixture +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 + + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(asdict(StartArgs()))) + get_env_start_args.cache_clear() + yield + get_env_start_args.cache_clear() + + @pytest.fixture def kv_allocator_environment(monkeypatch: pytest.MonkeyPatch): """Provide a service identity and reclaim the allocator's real shared counters.""" diff --git a/unit_tests/distributed/test_ordered_tp_sum.py b/unit_tests/distributed/test_ordered_tp_sum.py new file mode 100644 index 0000000000..bcf4885da4 --- /dev/null +++ b/unit_tests/distributed/test_ordered_tp_sum.py @@ -0,0 +1,110 @@ +import datetime +import json +import os +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + + +def _worker(rank, world, rendezvous): + torch.cuda.set_device(rank) + from lightllm.server.core.objs.start_args_type import StartArgs + + settings = StartArgs() + settings.hardware_platform, settings.data_type = "cuda", "bf16" + os.environ["LIGHTLLM_START_ARGS"] = json.dumps(vars(settings)) + for name, value in { + "LIGHTLLM_DP_WORLD_SIZE": world, + "LIGHTLLM_GLOBAL_WORLD_SIZE": world, + "LIGHTLLM_NODE_WORLD_SIZE": world, + "LIGHTLLM_CURRENT_DEVICE_ID": rank, + "LIGHTLLM_CURRENT_RANK_IN_DP": rank, + "LIGHTLLM_CURRENT_RANK_IN_NODE": rank, + "LIGHTLLM_GLOBAL_RANK": rank, + "LIGHTLLM_DP_SIZE": 1, + "LIGHTLLM_GLOBAL_DP_RANK": 0, + "LIGHTLLM_DP_RANK_IN_NODE": 0, + }.items(): + os.environ[name] = str(value) + dist.init_process_group( + "nccl", + init_method="file://" + rendezvous, + rank=rank, + world_size=world, + timeout=datetime.timedelta(seconds=180), + ) + from lightllm.distributed.communication_op import CustomProcessGroup + from lightllm.distributed.symm_mem_all_reduce import SymmMemAllreduce + from lightllm.common.basemodel.tp_reduce import ordered_tp_sum + + try: + symmetric = SymmMemAllreduce(dist.group.WORLD, rank) + available = torch.tensor(int(not symmetric.disabled), device=rank) + dist.all_reduce(available, op=dist.ReduceOp.MIN) + if not available.item(): + if rank == 0: + Path(rendezvous + ".unavailable").touch() + return + group = CustomProcessGroup.__new__(CustomProcessGroup) + group.device_group, group.dp_world_size = dist.group.WORLD, world + group.symm_mem_reduce, group.flashinfer_reduce = symmetric, None + for rows in (1, 8, 24, 512): + source = torch.empty((rows, 2560), dtype=torch.bfloat16, device=rank) + reference = torch.empty((world * rows, 2560), dtype=source.dtype, device=rank) + scratch = torch.empty_like(source) + + def expected(): + dist.all_gather_into_tensor(reference, source) + return reference.view(world, rows, 2560).double().sum(0).to(source.dtype) + + for iteration in range(4): + if world >= 4 and iteration == 0: + source.fill_((2 ** 40, 1, -(2 ** 40), 1)[rank] if rank < 4 else 0) + else: + source.copy_((torch.arange(source.numel(), device=rank).reshape_as(source) % 257 + rank) / 64) + source.add_(iteration / 8) + gold = expected() + gathered = symmetric.ordered_all_gather(source) + assert gathered is not None + scratch.fill_(rank + 2) + symmetric.all_reduce(scratch) + torch.testing.assert_close(gathered.view_as(reference), reference, rtol=0, atol=0) + actual = ordered_tp_sum(source.clone(), group, world) + torch.testing.assert_close(actual, gold, rtol=0, atol=0) + if rows % world == 0: + actual = ordered_tp_sum(source, group, world, scatter_rank=rank) + torch.testing.assert_close(actual, gold.chunk(world)[rank], rtol=0, atol=0) + + output = torch.empty_like(source) + graph = torch.cuda.CUDAGraph() + torch.cuda.synchronize() + with torch.cuda.graph(graph): + output.copy_(source) + ordered_tp_sum(output, group, world) + for iteration in range(8): + source.fill_(rank + iteration / 16) + gold = expected() + graph.replay() + torch.testing.assert_close(output, gold, rtol=0, atol=0) + graph.reset() + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("world", [2, 8]) +def test_ordered_sum_preserves_shards_across_collectives_and_graph_replays(world, tmp_path): + if torch.cuda.device_count() < world: + pytest.skip(f"Requires {world} CUDA GPUs") + if any(torch.cuda.get_device_capability(rank)[0] < 9 for rank in range(world)): + pytest.skip("Requires Hopper or newer GPUs with multicast support") + from lightllm.utils.device_utils import has_nvlink + + if not has_nvlink(): + pytest.skip("Requires an NVLink-connected node") + rendezvous = str(tmp_path / "rendezvous") + mp.spawn(_worker, args=(world, rendezvous), nprocs=world, join=True) + if Path(rendezvous + ".unavailable").exists(): + pytest.skip("Symmetric multicast memory is unavailable") diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 530b95436a..f583bcf7da 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -48,11 +48,13 @@ def test_nope_cache_config_uses_native_mla_width(): @pytest.mark.parametrize("small_page", [False, True]) @pytest.mark.parametrize("tp_world_size", [1, 4]) @pytest.mark.parametrize("mtp_step", [0, 2]) +@pytest.mark.parametrize("mode", ["native", "replay"]) def test_hybrid_checkpoint_restore_and_packed_kv_copy( monkeypatch, small_page, tp_world_size, mtp_step, + mode, ): monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") @@ -106,6 +108,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( config=config, request_capacity=4, mtp_step=mtp_step, + ssm_state_mode=mode, full_state=main_full, full_allocator=FullAttKVAllocator(full.capacity, full.page_size), draft_full_state=draft_full, @@ -143,6 +146,16 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( linear.ssm_state[:, 0].normal_() conv = linear.conv_state[:, 0, ..., :3].clone() ssm = linear.ssm_state[:, 0].clone() + if mode == "replay": + linear.raw_keys.zero_() + linear.raw_values.zero_() + linear.gates.fill_(-0.5) + linear.betas.zero_() + linear.cursors[0] = 1 + ssm = linear.accepted_ssm(0).clone() + linear.cursors[2] = 3 + if mtp_step: + linear.mtp_state_index[2] = mtp_step resources.save_bounded_att_runtime_to_pages( req_indices=[0], state_token_lens=[8], @@ -162,12 +175,20 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( ) torch.cuda.synchronize() assert torch.equal(linear.conv_state[:, 2, ..., :3], conv) - assert torch.equal(linear.ssm_state[:, 2 * (mtp_step + 1)], ssm) + restored_slot = 2 if mode == "replay" else 2 * (linear.mtp_step + 1) + assert torch.equal(linear.ssm_state[:, restored_slot], ssm) + if mode == "replay": + assert linear.cursors[2] == 0 + if mtp_step: + assert linear.mtp_state_index[2] == 0 # 整组边界不依赖旧窗口，恢复大小页无需改动无重叠窗口。 assert torch.equal(tail.buffer, unchanged) resources.reset_bounded_att_state([2]) assert not linear.conv_state[:, 2].any() - assert not linear.ssm_state[:, 2 * (mtp_step + 1) : 3 * (mtp_step + 1)].any() + if mode == "replay": + assert not linear.ssm_state[:, 2].any() + else: + assert not linear.ssm_state[:, 2 * (linear.mtp_step + 1) : 3 * (linear.mtp_step + 1)].any() assert not tail.buffer[:, 2, :, : tail.config.candidate_dim].any() assert torch.isneginf(tail.buffer[:, 2, :, tail.config.candidate_dim :]).all() assert torch.equal(tail.buffer[:, [0, 1, 3]], unchanged[:, [0, 1, 3]]) diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 8a1d70e3ac..bc6c1428cb 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -5,7 +5,9 @@ from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, KDAState from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import ( + fused_recurrent_kda, +) from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools @@ -23,7 +25,11 @@ def test_kda_mtp_matches_single_token_decode_and_restores_accepted_state(mtp_ste states = torch.randn(6 * width, heads, dim, dim, device="cuda", dtype=state_dtype) * 0.1 accepted = torch.tensor([width, 1, 1], device="cuda", dtype=torch.int32) for lengths in ([width, 1, 0], [2, width, 0], [1, 1, 0]): - cu = torch.tensor([0, lengths[0], sum(lengths), sum(lengths)], device="cuda", dtype=torch.int32) + cu = torch.tensor( + [0, lengths[0], sum(lengths), sum(lengths)], + device="cuda", + dtype=torch.int32, + ) tokens = sum(lengths) q, k, v, gate = [torch.randn(1, tokens, heads, dim, device="cuda", dtype=torch.bfloat16) for _ in range(4)] beta = torch.randn(1, tokens, heads, device="cuda", dtype=torch.bfloat16) @@ -100,7 +106,12 @@ def prefill(values, destination, tail): prefill(raw[:prefix], packed, ring) # Alternate expanded verification, rewound single-token drafting, and re-verification. - for start, count in [(prefix, width), (prefix + 1, 1), (prefix + 2, 1), (prefix + 1, width)]: + for start, count in [ + (prefix, width), + (prefix + 1, 1), + (prefix + 2, 1), + (prefix + 1, width), + ]: raw[start : start + count] = torch.randn(count, 256, device="cuda", dtype=torch.bfloat16) lengths = torch.arange(start + 1, start + count + 1, device="cuda", dtype=torch.int32) compress_pools( @@ -117,7 +128,11 @@ def prefill(values, destination, tail): b_mtp_index=torch.arange(count, device="cuda", dtype=torch.int32), ) expected = torch.zeros_like(packed) - prefill(raw[: start + count], expected, torch.zeros(2, 4, 256, device="cuda", dtype=torch.bfloat16)) + prefill( + raw[: start + count], + expected, + torch.zeros(2, 4, 256, device="cuda", dtype=torch.bfloat16), + ) closing = torch.arange(start, start + count, device="cuda") closing = closing[(closing + 1) % 4 == 0] assert torch.equal(packed[closing], expected[closing]) @@ -125,17 +140,26 @@ def prefill(values, destination, tail): @pytest.mark.parametrize("dynamic", [False, True]) @pytest.mark.parametrize("role", ["main", "draft0"]) -def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): +@pytest.mark.parametrize("mode,mtp_step", [("native", 2), ("replay", 2), ("native", 0), ("replay", 0)]) +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.triton_kernel.linear_att.causal_conv1d import causal_conv1d_update + 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, + ) torch.manual_seed(53) - heads, dim, width = 2, 128, 3 + heads, dim, width = 2, 128, mtp_step + 1 hidden = heads * dim - lengths = [2, 1, 3] if dynamic else [3, 3, 3] + lengths = [2, 1, 3] if dynamic and mtp_step else [width] * 3 requests = [2, 0, 1] request_rows = [req for req, length in zip(requests, lengths) for _ in range(length)] offsets = [offset for length in lengths for offset in range(length)] @@ -143,10 +167,10 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): request_rows += [3, 3] offsets += [0, 0] tokens = len(request_rows) - conv = torch.randn(4, 3 * hidden, 5, device="cuda", dtype=torch.bfloat16) + conv = torch.randn(4, 3 * hidden, 3 + mtp_step, device="cuda", dtype=torch.bfloat16) ssm = torch.randn(4 * width, heads, dim, dim, device="cuda") * 0.01 old_conv, expected_ssm = conv.clone(), ssm.clone() - accepted_offsets = torch.tensor([1, 2, 0, 0], dtype=torch.int32, device="cuda") + accepted_offsets = torch.tensor([1, 2, 0, 0] if mtp_step else [0, 0, 0, 0], dtype=torch.int32, device="cuda") mixed = torch.randn(tokens, 3 * hidden, device="cuda", dtype=torch.bfloat16) gate = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) beta = torch.randn(tokens, heads, device="cuda", dtype=torch.bfloat16) @@ -184,7 +208,7 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): expected_ssm[req * width + token] = state[0] row += 1 backend = SimpleNamespace( - mtp_step=2, + mtp_step=mtp_step, tp_num_heads=heads, head_dim=dim, tp_hidden_size=hidden, @@ -203,11 +227,20 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): conv_dtype=conv.dtype, ssm_dtype=ssm.dtype, ) + selected = None for model_role in ("main", "draft0"): - linear = KDAState(config, request_capacity=4, mtp_step=2, device="meta") - linear.mtp_state_index = accepted_offsets if model_role == role else torch.zeros_like(accepted_offsets) + state_cls = ReplayKDAState if mode == "replay" else KDAState + linear = state_cls(config, request_capacity=4, mtp_step=mtp_step, device="cuda") + if mtp_step: + linear.mtp_state_index.copy_(accepted_offsets if model_role == role else torch.zeros_like(accepted_offsets)) if model_role == role: - linear.conv_state, linear.ssm_state = conv.unsqueeze(0), ssm.unsqueeze(0) + linear.conv_state[0].copy_(conv) + if mode == "native": + linear.ssm_state[0].copy_(ssm) + else: + indexes = torch.arange(4, device="cuda") * width + accepted_offsets + linear.ssm_state[0].copy_(ssm[indexes]) + selected = linear resources.register_state(model_role, linear) resources.finish_state_registration() manager = SimpleNamespace(HOLD_REQUEST_ID=3) @@ -223,7 +256,8 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): linear_A_log=SimpleNamespace(weight=a), linear_dt_bias=SimpleNamespace(weight=bias), ) - state = KDADecodeAttState(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) + 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, @@ -232,11 +266,40 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): 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 + layer_weight=weight, + layer_num=0, + mixed_qkv=mixed, + raw_gate=gate, + raw_beta=beta, ), ), ) - torch.testing.assert_close(actual.reshape(tokens, heads, dim)[:row], torch.cat(expected), atol=2e-3, rtol=1e-2) - torch.testing.assert_close(ssm, expected_ssm, atol=2e-3, rtol=1e-2) - if dynamic: - assert torch.equal(conv[3], old_conv[3]) + 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, + ) + else: + ids = torch.tensor(requests, dtype=torch.int32, device="cuda") + if mtp_step: + accepted = torch.zeros(4, dtype=torch.int32, device="cuda") + accepted[ids] = torch.tensor(lengths, dtype=torch.int32, device="cuda") - 1 + selected.accept_updates(ids, accepted) + for req, length in zip(requests, lengths): + torch.testing.assert_close( + selected.accepted_ssm(req)[0], + expected_ssm[req * width + length - 1], + atol=2e-3, + rtol=1e-2, + ) + if dynamic and mtp_step: + assert torch.equal(selected.conv_state[0, 3], old_conv[3]) diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index e72e4d6bee..7dbe235831 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -46,7 +46,7 @@ def make_mems(monkeypatch, request): get_env_start_args.cache_clear() set_env_start_args(dataclasses.asdict(args)) - def create(tp, target_mtp_step=None): + def create(tp, target_mtp_step=None, mode="native"): local_mtp_step = mtp_step if target_mtp_step is None else target_mtp_step config = Glm5NextResourceConfig( tp_world_size=tp, @@ -87,6 +87,7 @@ def create(tp, target_mtp_step=None): config=config, request_capacity=4, mtp_step=local_mtp_step, + ssm_state_mode=mode, tp_rank=rank, full_state=main, full_allocator=FullAttKVAllocator(full.capacity, full.page_size), @@ -112,16 +113,21 @@ def create(tp, target_mtp_step=None): @pytest.mark.parametrize("prefill_tp,decode_tp", [(1, 1), (4, 4), (1, 4), (4, 1)]) @pytest.mark.parametrize("remainder", range(4)) @pytest.mark.parametrize("decode_mtp_step", [0, 1, 5]) +@pytest.mark.parametrize("source_mode", ["native", "replay"]) +@pytest.mark.parametrize("dest_mode", ["native", "replay"]) def test_pd_kv_and_runtime_state_roundtrip( make_mems, prefill_tp, decode_tp, remainder, decode_mtp_step, + source_mode, + dest_mode, ): - source, dest = make_mems(prefill_tp), make_mems(decode_tp, target_mtp_step=decode_mtp_step) + source = make_mems(prefill_tp, mode=source_mode) + dest = make_mems(decode_tp, target_mtp_step=decode_mtp_step, mode=dest_mode) + dst_ssm_stride = 1 if dest_mode == "replay" else dest[0].linear.mtp_step + 1 mtp_size = get_env_start_args().mtp_step + 1 - dst_mtp_size = decode_mtp_step + 1 layers = sum(len(state.config.logical_layer_groups) for state in source[0].full_states) src_req, dst_req = 0, 2 length = 8 + remainder @@ -159,7 +165,16 @@ def test_pd_kv_and_runtime_state_roundtrip( state.storage.view(torch.uint8)[:, src_indexes] = packed[offset:end] offset = end mem.linear.conv_state[:, src_req, ..., :3].copy_(conv[:, :, heads].reshape(6, -1, 3)) - mem.linear.ssm_state[:, src_req * mtp_size].copy_(ssm[:, heads]) + src_slot = src_req if source_mode == "replay" else src_req * (mem.linear.mtp_step + 1) + mem.linear.ssm_state[:, src_slot].copy_(ssm[:, heads]) + if source_mode == "replay": + backend = mem.linear + backend.raw_keys.zero_() + backend.raw_values.zero_() + backend.gates.fill_(-0.5) + backend.betas.zero_() + backend.cursors[src_req] = 1 + ssm[:, heads].copy_(mem.linear.accepted_ssm(src_req)) offset = 0 for state in mem.tails: end = offset + len(state.config.logical_layer_ids) @@ -173,6 +188,10 @@ def test_pd_kv_and_runtime_state_roundtrip( state.storage.zero_() mem.linear.conv_state.normal_() mem.linear.ssm_state.normal_() + if dest_mode == "replay": + mem.linear.cursors[dst_req] = 3 + if mem.linear.mtp_state_index is not None: + mem.linear.mtp_state_index[dst_req] = decode_mtp_step for state in mem.tails: state.buffer.normal_() original_tails.append(torch.cat([state.buffer[:, dst_req] for state in mem.tails]).clone()) @@ -181,7 +200,7 @@ def test_pd_kv_and_runtime_state_roundtrip( x[:, torch.tensor([0, 1, 3], device="cuda") * stride].clone() for x, stride in ( (mem.linear.conv_state, 1), - (mem.linear.ssm_state, dst_mtp_size), + (mem.linear.ssm_state, dst_ssm_stride), *((state.buffer, 1) for state in mem.tails), ) ] @@ -235,18 +254,20 @@ def test_pd_kv_and_runtime_state_roundtrip( source_field.tensor[:, src_indexes].contiguous().view(torch.uint8), ) assert torch.equal(mem.linear.conv_state[:, dst_req, ..., :3], conv[:, :, heads].reshape(6, -1, 3)) - assert torch.equal(mem.linear.ssm_state[:, dst_req * dst_mtp_size], ssm[:, heads]) + assert torch.equal(mem.linear.ssm_state[:, dst_req * dst_ssm_stride], ssm[:, heads]) restored_tail = torch.cat([state.buffer[:, dst_req] for state in mem.tails]) expected_tail = original_tails[rank].clone() for position in range(length - 4, length): expected_tail[:, position % (4 + decode_mtp_step)] = tail[:, position % tail.shape[1]] assert torch.equal(restored_tail, expected_tail) + if dest_mode == "replay": + assert mem.linear.cursors[dst_req] == 0 if mem.linear.mtp_state_index is not None: assert mem.linear.mtp_state_index[dst_req] == 0 for (state, stride), expected in zip( ( (mem.linear.conv_state, 1), - (mem.linear.ssm_state, dst_mtp_size), + (mem.linear.ssm_state, dst_ssm_stride), *((state.buffer, 1) for state in mem.tails), ), untouched[rank], diff --git a/unit_tests/models/qwen2_vl/test_rotary_pos_emb.py b/unit_tests/models/qwen2_vl/test_rotary_pos_emb.py index 34e41578a6..98f770cd77 100644 --- a/unit_tests/models/qwen2_vl/test_rotary_pos_emb.py +++ b/unit_tests/models/qwen2_vl/test_rotary_pos_emb.py @@ -1,4 +1,3 @@ -import math import torch import pytest @@ -51,5 +50,21 @@ def test_triton_matches_reference(shape): assert torch.allclose(out, ref, rtol=1e-2, atol=1e-2), "Triton 与参考实现不一致" +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_strided_qkv_view_uses_output_strides(dtype): + torch.manual_seed(1907) + qkv = torch.randn(37, 3, 3, 72, device="cuda", dtype=dtype) + tensor = qkv[:, 1] + assert not tensor.is_contiguous() + original = qkv.clone() + angles = torch.randn(37, 36, device="cuda", dtype=torch.float32) + cos, sin = angles.cos(), angles.sin() + expected = apply_rotary_pos_emb_vision(tensor.unsqueeze(0), cos, sin).squeeze(0) + actual = apply_rotary_pos_emb_triton(tensor, cos, sin) + tolerance = 0.016 if dtype == torch.bfloat16 else 2e-5 + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + torch.testing.assert_close(qkv, original, atol=0, rtol=0) + + if __name__ == "__main__": pytest.main() diff --git a/unit_tests/models/qwen3next/test_transformer_layer_infer_strided_gate.py b/unit_tests/models/qwen3next/test_transformer_layer_infer_strided_gate.py new file mode 100644 index 0000000000..5bfae4efa7 --- /dev/null +++ b/unit_tests/models/qwen3next/test_transformer_layer_infer_strided_gate.py @@ -0,0 +1,33 @@ +from types import SimpleNamespace + +import torch + +from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import Qwen3NextTransformerLayerInfer + + +def test_linear_post_passes_strided_gate_without_materializing(): + tokens, heads, head_dim = 3, 4, 8 + packed_dim = heads * head_dim + 16 + packed = torch.randn((tokens, packed_dim)) + z = packed[:, 8 : 8 + heads * head_dim].view(tokens, heads, head_dim) + core_attn_out = torch.randn((1, tokens, heads, head_dim)) + calls = [] + + def linear_norm(input, gate_value, eps): + calls.append((input, gate_value, eps)) + return input + + layer_weight = SimpleNamespace( + linear_norm=linear_norm, + linear_out_proj=SimpleNamespace(mm=lambda input: input), + ) + layer_infer = object.__new__(Qwen3NextTransformerLayerInfer) + layer_infer.eps_ = 1e-6 + + output = layer_infer._linear_post(core_attn_out, z, layer_weight) + + assert len(calls) == 1 + assert calls[0][0].shape == (tokens * heads, head_dim) + assert calls[0][1] is z + assert not calls[0][1].is_contiguous() + assert output.shape == (tokens, heads * head_dim) diff --git a/unit_tests/models/qwen4_exp/__init__.py b/unit_tests/models/qwen4_exp/__init__.py new file mode 100644 index 0000000000..2ae28399f5 --- /dev/null +++ b/unit_tests/models/qwen4_exp/__init__.py @@ -0,0 +1 @@ +pass diff --git a/unit_tests/models/qwen4_exp/conftest.py b/unit_tests/models/qwen4_exp/conftest.py new file mode 100644 index 0000000000..074a35b5a1 --- /dev/null +++ b/unit_tests/models/qwen4_exp/conftest.py @@ -0,0 +1,84 @@ +import json + +import pytest + +from .runtime import runtime # noqa: F401 -- shared pytest fixture + +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils import envs_utils + + +@pytest.fixture(autouse=True) +def native_runtime(monkeypatch): + args = StartArgs() + args.hardware_platform = "cuda" + args.mtp_step = 0 + args.enable_ep_moe = False + args.auto_update_redundancy_expert = False + args.enable_tpsp_mix_mode = False + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(vars(args))) + for name, value in { + "LIGHTLLM_DP_WORLD_SIZE": 1, + "LIGHTLLM_CURRENT_RANK_IN_DP": 0, + "LIGHTLLM_CURRENT_DEVICE_ID": 0, + "LIGHTLLM_GLOBAL_RANK": 0, + "LIGHTLLM_GLOBAL_WORLD_SIZE": 1, + }.items(): + monkeypatch.setenv(name, str(value)) + envs_utils.get_env_start_args.cache_clear() + envs_utils.get_added_mtp_kv_layer_num.cache_clear() + envs_utils.get_redundancy_expert_num.cache_clear() + envs_utils.get_redundancy_expert_ids.cache_clear() + yield + envs_utils.get_env_start_args.cache_clear() + envs_utils.get_added_mtp_kv_layer_num.cache_clear() + envs_utils.get_redundancy_expert_num.cache_clear() + envs_utils.get_redundancy_expert_ids.cache_clear() + + +@pytest.fixture +def config(): + return dict( + hidden_size=32, + n_embed=32, + num_hidden_layers=4, + n_layer=4, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + full_attention_interval=4, + layer_types=["linear_attention", "linear_attention", "linear_attention", "full_attention"], + hc_count=2, + hc_lowrank=8, + rms_norm_eps=1e-6, + hidden_act="silu", + output_gate_type="sigmoid", + linear_num_key_heads=1, + linear_num_value_heads=2, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_conv_kernel_dim=4, + num_experts=2, + num_experts_per_tok=1, + norm_topk_prob=True, + moe_intermediate_size=32, + shared_expert_intermediate_size=32, + indexer_n_heads=2, + indexer_kv_heads=1, + indexer_head_dim=16, + indexer_budget=8, + indexer_compress_ratio=4, + ple_layer_ids=[2], + ple_embed_dim=32, + ple_conv_kernel_size=4, + ngram_size=3, + heads_per_ngram=2, + ngram_vocab_size_base=5, + make_ngram_vocab_size_divisible_by=4, + split_ngram_parts=4, + eos_token_id=63, + vocab_size=64, + partial_rotary_factor=0.5, + rope_scaling={"mrope_section": [1, 1, 0]}, + tie_word_embeddings=False, + ) diff --git a/unit_tests/models/qwen4_exp/helpers.py b/unit_tests/models/qwen4_exp/helpers.py new file mode 100644 index 0000000000..462a5863ea --- /dev/null +++ b/unit_tests/models/qwen4_exp/helpers.py @@ -0,0 +1,133 @@ +import torch + + +def rand(*shape): + return torch.randn(*shape, dtype=torch.bfloat16) * 0.1 + + +def add_hyper(weights, prefix, config, combine=True): + hidden, branches, rank = config["hidden_size"], config["hc_count"], config["hc_lowrank"] + width = hidden * branches + weights[prefix + ".hc_norm.weight"] = rand(width) + weights[prefix + ".input_mix_weight_down.weight"] = rand(rank, width) + weights[prefix + ".input_mix_weight_up.weight"] = rand(width, rank) + if combine: + weights[prefix + ".block_inject_weight.weight"] = rand(branches, width) + + +def checkpoint(config, layer): + hidden, inter, experts = config["hidden_size"], config["moe_intermediate_size"], config["num_experts"] + prefix = f"model.language_model.layers.{layer}" + weights = {} + for block in ["attn_hyper_connection", "mlp_hyper_connection"]: + add_hyper(weights, prefix + "." + block, config) + weights[prefix + ".mlp.gate.weight"] = rand(experts, hidden) + weights[prefix + ".mlp.experts.gate_up_proj"] = rand(experts, 2 * inter, hidden) + weights[prefix + ".mlp.experts.down_proj"] = rand(experts, hidden, inter) + weights[prefix + ".mlp.shared_expert_gate.weight"] = rand(1, hidden) + for proj, shape in [("gate", (inter, hidden)), ("up", (inter, hidden)), ("down", (hidden, inter))]: + weights[prefix + f".mlp.shared_expert.{proj}_proj.weight"] = rand(*shape) + if config["layer_types"][layer] == "linear_attention": + keydim = config["linear_num_key_heads"] * config["linear_key_head_dim"] + vdim = config["linear_num_value_heads"] * config["linear_value_head_dim"] + channels = 2 * keydim + vdim + weights[prefix + ".linear_attn.conv1d.weight"] = rand(channels, 1, 4) + for proj, rows in [ + ("qkv", channels), + ("z", vdim), + ("b", config["linear_num_value_heads"]), + ("a", config["linear_num_value_heads"]), + ]: + weights[prefix + f".linear_attn.in_proj_{proj}.weight"] = rand(rows, hidden) + weights[prefix + ".linear_attn.out_proj.weight"] = rand(hidden, vdim) + weights[prefix + ".linear_attn.A_log"] = torch.zeros(config["linear_num_value_heads"], dtype=torch.float32) + weights[prefix + ".linear_attn.dt_bias"] = torch.ones(config["linear_num_value_heads"], dtype=torch.float32) + weights[prefix + ".linear_attn.norm.weight"] = 1 + rand(config["linear_value_head_dim"]) + else: + dim, qheads, kvheads = config["head_dim"], config["num_attention_heads"], config["num_key_value_heads"] + for name, shape in [ + ("q_proj", (qheads * dim * 2, hidden)), + ("k_proj", (kvheads * dim, hidden)), + ("v_proj", (kvheads * dim, hidden)), + ("o_proj", (hidden, qheads * dim)), + ]: + weights[prefix + f".self_attn.{name}.weight"] = rand(*shape) + for name in ["q_norm", "k_norm"]: + weights[prefix + f".self_attn.{name}.weight"] = rand(dim) + weights[prefix + ".self_attn.indexer.index_qk_proj.weight"] = rand( + (config["indexer_n_heads"] + 1) * config["indexer_head_dim"], hidden + ) + for name in ["q_layernorm", "k_layernorm"]: + weights[prefix + f".self_attn.indexer.{name}.weight"] = rand(config["indexer_head_dim"]) + if layer + 1 in config["ple_layer_ids"]: + width, embed = config["hc_count"] * hidden, config["ple_embed_dim"] + weights[prefix + ".ple.key_proj.weight"] = rand(width, embed) + weights[prefix + ".ple.value_proj.weight"] = rand(hidden, embed) + weights[prefix + ".ple.conv1d.weight"] = rand(width, 1, 4) + for name in ["norm_key", "norm_query", "norm_conv"]: + weights[prefix + f".ple.{name}.weight"] = rand(width) + weights[prefix + ".ple.ple_embedding.layer_multipliers"] = torch.tensor( + [23703573157769, 20109073645365, 8052911324071], dtype=torch.int64 + ) + weights[prefix + ".ple.ple_embedding.ngram_heads_vocab_sizes"] = torch.tensor([5, 7, 11, 13], dtype=torch.int64) + weights[prefix + ".ple.ple_embedding.ngram_heads_offsets"] = torch.tensor([0, 5, 12, 23], dtype=torch.int64) + table = rand(36, 8) + for shard, values in enumerate(table.chunk(4)): + weights[prefix + f".ple.ple_embedding.ngram_embedding.shard_{shard}.weight"] = values + return weights + + +def prepost_checkpoint(config): + weights = { + "model.language_model.embed_tokens.weight": rand(config["vocab_size"], config["hidden_size"]), + "lm_head.weight": rand(config["vocab_size"], config["hidden_size"]), + } + add_hyper(weights, "model.language_model.hyper_connection_mixer", config, False) + return weights + + +def reference_hyper(x, weights, prefix, config, combine=True): + grouped = x.float().unflatten(-1, (config["hc_count"], config["hidden_size"])) + normalized = grouped * torch.rsqrt(grouped.square().mean(-1, keepdim=True) + config["rms_norm_eps"]) + normalized = (normalized.flatten(-2) * (1 + weights[prefix + ".hc_norm.weight"].cuda().float())).to(x.dtype) + low = torch.nn.functional.silu( + torch.nn.functional.linear(normalized, weights[prefix + ".input_mix_weight_down.weight"].cuda()) + / config["hc_count"] + ) + gates = torch.nn.functional.linear(low, weights[prefix + ".input_mix_weight_up.weight"].cuda()).sigmoid() + mixed = (normalized * gates).unflatten(-1, (config["hc_count"], config["hidden_size"])).mean(-2) + if not combine: + return mixed + injection = 2 * torch.sigmoid( + torch.nn.functional.linear(normalized, weights[prefix + ".block_inject_weight.weight"].cuda()) + / config["hc_count"] + ) + return mixed, injection + + +def reference_moe(x, weights, prefix, config): + logits = torch.nn.functional.linear(x, weights[prefix + ".mlp.gate.weight"].cuda()) + prob = logits.float().softmax(-1) + values, indices = prob.topk(config["num_experts_per_tok"], dim=-1) + if config["norm_topk_prob"]: + values /= values.sum(-1, keepdim=True) + values = values.to(x.dtype) + out = torch.zeros_like(x) + inter = config["moe_intermediate_size"] + + for expert in indices.unique(sorted=True).tolist(): + top_position, token = torch.where((indices == expert).T) + gu = torch.nn.functional.linear(x[token], weights[prefix + ".mlp.experts.gate_up_proj"][expert].cuda()) + output = torch.nn.functional.linear( + torch.nn.functional.silu(gu[:, :inter]) * gu[:, inter:], + weights[prefix + ".mlp.experts.down_proj"][expert].cuda(), + ) + weighted = output * values[token, top_position, None] + out.index_add_(0, token, weighted.to(out.dtype)) + gate = torch.nn.functional.linear(x, weights[prefix + ".mlp.shared_expert_gate.weight"].cuda()).sigmoid() + up = torch.nn.functional.linear(x, weights[prefix + ".mlp.shared_expert.up_proj.weight"].cuda()) + g = torch.nn.functional.linear(x, weights[prefix + ".mlp.shared_expert.gate_proj.weight"].cuda()) + shared = torch.nn.functional.linear( + torch.nn.functional.silu(g) * up, weights[prefix + ".mlp.shared_expert.down_proj.weight"].cuda() + ) + return out + gate * shared diff --git a/unit_tests/models/qwen4_exp/lexical_helpers.py b/unit_tests/models/qwen4_exp/lexical_helpers.py new file mode 100644 index 0000000000..ebd26a3389 --- /dev/null +++ b/unit_tests/models/qwen4_exp/lexical_helpers.py @@ -0,0 +1,32 @@ +import torch + +from lightllm.models.qwen4_exp.layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight +from lightllm.models.qwen4_exp.batch import PackedBatchMetadata +from lightllm.models.qwen4_exp.lexical_cache import LexicalRequestCache +from lightllm.common.quantization import Quantcfg +from .helpers import checkpoint +from .reference import reference_ple + + +def make_cache(config): + source = checkpoint(config, 1) + weights = Qwen4ExpTransformerLayerWeight(1, torch.bfloat16, config, Quantcfg(config)) + weights.load_hf_weights(dict(source)) + weights.verify_load() + return LexicalRequestCache(weights, 8), source + + +def metadata(requests, lengths, queries): + tensors = [torch.tensor(values, device="cuda", dtype=torch.int32) for values in (requests, lengths, queries)] + return PackedBatchMetadata.prepare(*tensors, sum(queries), 7, 4) + + +def reference_segment(residual, tokens, source, config, ngram, convolution): + updates, states = [], [(ngram.clone(), convolution.clone())] + for row in range(tokens.numel()): + update, ngram, convolution = reference_ple( + residual[row : row + 1], tokens[row : row + 1], source, config, ngram, convolution + ) + updates.append(update) + states.append((ngram.clone(), convolution.clone())) + return torch.cat(updates), states diff --git a/unit_tests/models/qwen4_exp/mtp_reference.py b/unit_tests/models/qwen4_exp/mtp_reference.py new file mode 100644 index 0000000000..e9d4332c42 --- /dev/null +++ b/unit_tests/models/qwen4_exp/mtp_reference.py @@ -0,0 +1,27 @@ +import torch.nn.functional as F + +from .reference import norm, reference_layer +from .helpers import reference_hyper + + +def reference_fusion(embedding, hidden, weights, config): + embedding = norm(embedding, weights["mtp.pre_fc_norm_embedding.weight"], config["rms_norm_eps"]) + hidden = norm(hidden, weights["mtp.pre_fc_norm_hidden.weight"], config["rms_norm_eps"]) + projected_embedding = F.linear(embedding, weights["mtp.fc_embedding.weight"].cuda()) + projected_hidden = F.linear( + hidden.view(-1, config["hc_count"], config["hidden_size"]), weights["mtp.fc_hidden.weight"].cuda() + ) + return (projected_hidden + projected_embedding[:, None]).flatten(1) + + +def reference_draft(tokens, hidden, weights, config, cos, sin, lengths): + embeddings = F.embedding(tokens, weights["model.language_model.embed_tokens.weight"].cuda()) + residual = reference_fusion(embeddings, hidden, weights, config) + mapped = { + name.replace("mtp.layers.0.", "model.language_model.layers.3.", 1): value + for name, value in weights.items() + if name.startswith("mtp.layers.0.") + } + residual = reference_layer(residual, mapped, config, cos, sin, lengths) + mixed = reference_hyper(residual, weights, "mtp.hyper_connection_mixer", config, combine=False) + return residual, F.linear(mixed, weights["lm_head.weight"].cuda()).float() diff --git a/unit_tests/models/qwen4_exp/reference.py b/unit_tests/models/qwen4_exp/reference.py new file mode 100644 index 0000000000..2557c479f5 --- /dev/null +++ b/unit_tests/models/qwen4_exp/reference.py @@ -0,0 +1,235 @@ +import torch +import torch.nn.functional as F + +from .helpers import reference_hyper, reference_moe + + +def grouped_norm(x, weight, config): + values = x.float().unflatten(-1, (config["hc_count"], config["hidden_size"])) + normalized = values * torch.rsqrt(values.square().mean(-1, keepdim=True) + config["rms_norm_eps"]) + return (normalized.flatten(-2) * (1 + weight.cuda().float())).to(x.dtype) + + +def reference_ple(residual, tokens, weights, config, ngram_history, conv_history): + prefix = "model.language_model.layers.1.ple" + history = torch.cat([ngram_history, tokens[:, None]], dim=1) + eos = config["eos_token_id"] + eos = eos[0] if isinstance(eos, list) else eos + context = config["ngram_size"] - 1 + shifted = [tokens] + for distance in range(1, context + 1): + preceding = history[:, context - distance : context] + shifted.append(torch.where((preceding == eos).any(-1), eos, history[:, context - distance])) + multipliers = weights[prefix + ".ple_embedding.layer_multipliers"].cuda() + sizes = weights[prefix + ".ple_embedding.ngram_heads_vocab_sizes"].cuda() + offsets = weights[prefix + ".ple_embedding.ngram_heads_offsets"].cuda() + indices = [] + for head in range(sizes.numel()): + order = head // config["heads_per_ngram"] + 2 + mixed = shifted[0] * multipliers[0] + for position in range(1, order): + mixed = mixed ^ (shifted[position] * multipliers[position]) + indices.append(mixed.remainder(sizes[head]) + offsets[head]) + embedding = reference_ngram_lookup(torch.stack(indices, -1), weights, prefix).flatten(-2) + branches, hidden = config["hc_count"], config["hidden_size"] + key = grouped_norm( + F.linear(embedding, weights[prefix + ".key_proj.weight"].cuda()), weights[prefix + ".norm_key.weight"], config + ).view(-1, branches, hidden) + query = grouped_norm(residual, weights[prefix + ".norm_query.weight"], config).view(-1, branches, hidden) + gate = (key * query).sum(-1, keepdim=True) / (hidden ** 0.5) + gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() + value = F.linear(embedding, weights[prefix + ".value_proj.weight"].cuda()) + gated = (gate.sigmoid() * value[:, None]).flatten(-2) + normalized = grouped_norm(gated, weights[prefix + ".norm_conv.weight"], config) + full = torch.cat([conv_history, normalized[:, None]], dim=1) + dilation = config["ngram_size"] + conv = F.conv1d( + full.transpose(1, 2), weights[prefix + ".conv1d.weight"].cuda(), groups=branches * hidden, dilation=dilation + ).squeeze(-1) + conv_context = (config["ple_conv_kernel_size"] - 1) * dilation + return ( + gated + F.silu(conv), + history[:, -context:].contiguous(), + full[:, -conv_context:].contiguous() if conv_context else full[:, :0].contiguous(), + ) + + +def reference_ngram_lookup(indices, weights, prefix): + name = prefix + ".ple_embedding.ngram_embedding" + first = weights[name + ".shard_0.weight"] + shard_rows, dim = first.shape + flat = indices.cpu().flatten() + shards = flat // shard_rows + result = torch.empty(flat.numel(), dim, dtype=first.dtype) + for shard in shards.unique().tolist(): + positions = torch.where(shards == shard)[0] + source = weights[name + f".shard_{shard}.weight"] + rows = (flat[positions] % shard_rows).to(source.device) + result[positions] = source[rows].cpu() + return result.to(indices.device).view(*indices.shape, dim) + + +def reference_gdn(mixed, weights, config, conv_history, recurrent): + prefix = "model.language_model.layers.1.linear_attn" + key_heads, value_heads = config["linear_num_key_heads"], config["linear_num_value_heads"] + key_dim, value_dim = config["linear_key_head_dim"], config["linear_value_head_dim"] + kd, vd = key_heads * key_dim, value_heads * value_dim + qkv = F.linear(mixed, weights[prefix + ".in_proj_qkv.weight"].cuda()) + full = torch.cat([conv_history, qkv[:, :, None]], dim=-1) + conv = (full.float() * weights[prefix + ".conv1d.weight"].squeeze(1).cuda().float()[None]).sum(-1).to(mixed.dtype) + conv = F.silu(conv) + q, k, v = torch.split(conv, [kd, kd, vd], dim=-1) + q, k = q.view(-1, key_heads, key_dim).float(), k.view(-1, key_heads, key_dim).float() + q = q * torch.rsqrt(q.square().sum(-1, keepdim=True) + 1e-6) + k = k * torch.rsqrt(k.square().sum(-1, keepdim=True) + 1e-6) + q = q.repeat_interleave(value_heads // key_heads, dim=1) / (key_dim ** 0.5) + k = k.repeat_interleave(value_heads // key_heads, dim=1) + v = v.view(-1, value_heads, value_dim).float() + a = F.linear(mixed, weights[prefix + ".in_proj_a.weight"].cuda()) + b = F.linear(mixed, weights[prefix + ".in_proj_b.weight"].cuda()) + decay = -weights[prefix + ".A_log"].cuda().float().exp() * F.softplus( + a.float() + weights[prefix + ".dt_bias"].cuda() + ) + beta = b.sigmoid().float() + next_state = recurrent * decay.exp()[:, :, None, None] + memory = (next_state * k[:, :, :, None]).sum(-2) + delta = (v - memory) * beta[:, :, None] + next_state = next_state + k[:, :, :, None] * delta[:, :, None, :] + core = (next_state * q[:, :, :, None]).sum(-2).to(mixed.dtype) + normed = core.float() * torch.rsqrt(core.float().square().mean(-1, keepdim=True) + 1e-6) + weighted = weights[prefix + ".norm.weight"].cuda() * normed.to(mixed.dtype) + z = F.linear(mixed, weights[prefix + ".in_proj_z.weight"].cuda()).view(-1, value_heads, value_dim) + gated = (weighted * z.float().sigmoid()).to(mixed.dtype) + output = F.linear(gated.flatten(1), weights[prefix + ".out_proj.weight"].cuda()) + return output, full[:, :, -3:].contiguous(), next_state + + +def reference_step(residual, tokens, source, config, ngram, ple, conv, recurrent): + lexical, ngram, ple = reference_ple(residual, tokens, source, config, ngram, ple) + expected = residual + lexical + mixed, gates = reference_hyper(expected, source, "model.language_model.layers.1.attn_hyper_connection", config) + output, conv, recurrent = reference_gdn(mixed, source, config, conv, recurrent) + expected = expected + (output[:, None] * gates[:, :, None]).flatten(1) + mixed, gates = reference_hyper(expected, source, "model.language_model.layers.1.mlp_hyper_connection", config) + output = reference_moe(mixed, source, "model.language_model.layers.1", config) + expected = expected + (output[:, None] * gates[:, :, None]).flatten(1) + return expected, ngram, ple, conv, recurrent + + +def norm(x, weight, eps=1e-6): + value = x.float() + return (value * torch.rsqrt(value.square().mean(-1, keepdim=True) + eps) * (1 + weight.cuda().float())).to(x.dtype) + + +def rotate(x, cos, sin): + width = cos.shape[-1] + a, b = x[..., :width].chunk(2, -1) + return torch.cat([x[..., :width] * cos + torch.cat([-b, a], -1) * sin, x[..., width:]], -1) + + +def reference_layer(residual, weights, config, cos, sin, lengths): + prefix = "model.language_model.layers.3" + mixed, gates = reference_hyper(residual, weights, prefix + ".attn_hyper_connection", config) + dim, heads, kvheads = config["head_dim"], config["num_attention_heads"], config["num_key_value_heads"] + qg = F.linear(mixed, weights[prefix + ".self_attn.q_proj.weight"].cuda()).view(-1, heads, dim * 2) + query, gate = qg.chunk(2, -1) + key = F.linear(mixed, weights[prefix + ".self_attn.k_proj.weight"].cuda()).view(-1, kvheads, dim) + value = F.linear(mixed, weights[prefix + ".self_attn.v_proj.weight"].cuda()).view(-1, kvheads, dim) + query = rotate(norm(query, weights[prefix + ".self_attn.q_norm.weight"]), cos[:, None], sin[:, None]) + key = rotate(norm(key, weights[prefix + ".self_attn.k_norm.weight"]), cos[:, None], sin[:, None]) + projected = F.linear(mixed, weights[prefix + ".self_attn.indexer.index_qk_proj.weight"].cuda()) + ih, idim = config["indexer_n_heads"], config["indexer_head_dim"] + iq = projected[:, : ih * idim].view(-1, ih, idim) + iq = rotate(norm(iq, weights[prefix + ".self_attn.indexer.q_layernorm.weight"]), cos[:, None], sin[:, None]) + raw = projected[:, ih * idim :] + ratio, budget = config["indexer_compress_ratio"], config["indexer_budget"] // config["indexer_compress_ratio"] + result = [] + start = 0 + for count in lengths: + for position in range(count): + complete = (position + 1) // ratio + selected = [] + if complete: + groups = raw[start : start + complete * ratio].view(complete, ratio, idim).float().mean(1).to(raw.dtype) + pooled = norm(groups, weights[prefix + ".self_attn.indexer.k_layernorm.weight"]) + pooled = rotate( + pooled, cos[start : start + complete * ratio : ratio], sin[start : start + complete * ratio : ratio] + ) + scores = (iq[start + position].float() @ pooled.float().T).relu().sum(0) / (idim ** 0.5) + blocks = scores.topk(min(complete, budget)).indices + selected = (blocks[:, None] * ratio + torch.arange(ratio, device=residual.device)).flatten().tolist() + selected += list(range(complete * ratio, position + 1)) + ids = torch.tensor(selected, device=residual.device, dtype=torch.int64) + start + out = F.scaled_dot_product_attention( + query[start + position].unsqueeze(0).unsqueeze(2), + key[ids].transpose(0, 1).unsqueeze(0), + value[ids].transpose(0, 1).unsqueeze(0), + enable_gqa=True, + ) + result.append(out[0, :, 0]) + start += count + attention = torch.stack(result) + attention = attention * gate.sigmoid() + output = F.linear(attention.flatten(1), weights[prefix + ".self_attn.o_proj.weight"].cuda()) + residual = residual + (output[:, None, :] * gates[:, :, None]).flatten(1) + mixed, gates = reference_hyper(residual, weights, prefix + ".mlp_hyper_connection", config) + update = reference_moe(mixed, weights, prefix, config) + return residual + (update[:, None, :] * gates[:, :, None]).flatten(1) + + +def reference_model(tokens, cos, sin, lengths, weights, config, input_embeddings=None): + if input_embeddings is None: + input_embeddings = F.embedding(tokens, weights["model.language_model.embed_tokens.weight"].cuda()) + residual = input_embeddings.repeat(1, config["hc_count"]) + for layer in range(config["n_layer"]): + prefix = f"model.language_model.layers.{layer}" + if config["layer_types"][layer] == "full_attention": + residual = reference_layer(residual, weights, config, cos, sin, lengths) + continue + if layer + 1 in config["ple_layer_ids"]: + updates, start = [], 0 + for length in lengths: + ngram = torch.full((1, 2), config["eos_token_id"], device="cuda", dtype=torch.int64) + convolution = torch.zeros(1, 9, 64, device="cuda", dtype=torch.bfloat16) + for row in range(start, start + length): + update, ngram, convolution = reference_ple( + residual[row : row + 1], tokens[row : row + 1], weights, config, ngram, convolution + ) + updates.append(update) + start += length + residual = residual + torch.cat(updates) + mixed, gates = reference_hyper(residual, weights, prefix + ".attn_hyper_connection", config) + + local = { + name.replace(prefix, "model.language_model.layers.1", 1): value + for name, value in weights.items() + if name.startswith(prefix + ".") + } + updates, start = [], 0 + for length in lengths: + channels = ( + 2 * config["linear_num_key_heads"] * config["linear_key_head_dim"] + + config["linear_num_value_heads"] * config["linear_value_head_dim"] + ) + convolution = torch.zeros( + 1, channels, config["linear_conv_kernel_dim"] - 1, device="cuda", dtype=torch.bfloat16 + ) + recurrent = torch.zeros( + 1, + config["linear_num_value_heads"], + config["linear_key_head_dim"], + config["linear_value_head_dim"], + device="cuda", + dtype=torch.float32, + ) + for row in range(start, start + length): + update, convolution, recurrent = reference_gdn( + mixed[row : row + 1], local, config, convolution, recurrent + ) + updates.append(update) + start += length + residual = residual + (torch.cat(updates)[:, None] * gates[:, :, None]).flatten(1) + mixed, gates = reference_hyper(residual, weights, prefix + ".mlp_hyper_connection", config) + residual = residual + (reference_moe(mixed, weights, prefix, config)[:, None] * gates[:, :, None]).flatten(1) + hidden = reference_hyper(residual, weights, "model.language_model.hyper_connection_mixer", config, combine=False) + return F.linear(hidden, weights["lm_head.weight"].cuda()).float() diff --git a/unit_tests/models/qwen4_exp/resource_helpers.py b/unit_tests/models/qwen4_exp/resource_helpers.py new file mode 100644 index 0000000000..0813aed3b8 --- /dev/null +++ b/unit_tests/models/qwen4_exp/resource_helpers.py @@ -0,0 +1,73 @@ +from types import SimpleNamespace + +import torch + +from lightllm.models.qwen4_exp.attention_state import Qwen4ExpAttentionState +from lightllm.models.qwen4_exp.lexical_state import Qwen4ExpLexicalState + + +class ResourceView: + def __init__(self, manager): + self.manager = manager + self.allocator = manager.full_att_allocator + self.size = 32 + (self.state,) = manager.get_full_attention_states("main") + self.qsa_states = { + layer: state.qsa_view(layer) + for state in manager.get_full_attention_states() + for group in state.config.logical_layer_groups + for layer in group + } + self.qsa_raw_keys = self.state.indexer.k_storage[:, :, 0] + self.qsa_pooled_keys = self.state.indexer.v_storage[:, :, 0] + self.qsa_key_cos = self.state.positions.k_storage[0, :, 0] + self.qsa_key_sin = self.state.positions.v_storage[0, :, 0] + + def alloc(self, count): + return self.allocator.alloc(count) + + def _slot_buffers(self): + return { + f"{state.state_name}.{name}": (tensor, 1) + for state in self.manager.get_full_attention_states() + for name, tensor in ( + ("k", state.k_storage), + ("v", state.v_storage), + ("raw", state.indexer.k_storage), + ("pooled", state.indexer.v_storage), + ("cos", state.positions.k_storage), + ("sin", state.positions.v_storage), + ) + } + + def get_index_kv_buffer(self, indices): + indices = torch.as_tensor(indices, dtype=torch.int64, device=self.state.device) + return {name: tensor.index_select(axis, indices) for name, (tensor, axis) in self._slot_buffers().items()} + + +def bind_states(model): + manager = model.att_state_manager + manager._get_states.cache_clear() + manager.finish_state_registration() + if manager.full_att_allocator is None: + manager.allocate_storage(32, 1, 8, 8, device="cuda") + else: + for state in manager.get_full_attention_states(): + if state.k_storage is None: + state.device = torch.device("cuda") + state.allocate_storage(32) + req = model.req_manager + req.bind_hold_page(manager.full_att_allocator.hold_token_indices) + lexical = next( + state for state in manager.get_bounded_attention_states("main") if isinstance(state, Qwen4ExpLexicalState) + ) + req.lexical_caches = lexical.lexical_caches + req.init_linear_att_state = lambda request: manager.reset_bounded_att_state([request.req_idx]) + req.snapshot_lexical = lambda request: { + layer: cache.snapshot(request) for layer, cache in lexical.lexical_caches.items() + } + req.restore_lexical = lambda request, snapshots: [ + cache.restore(request, snapshots[layer]) for layer, cache in lexical.lexical_caches.items() + ] + req.get_mamba_cache = lambda layer: manager.get_state("main", "gdn", layer).get_layer_view(layer) + return ResourceView(manager) diff --git a/unit_tests/models/qwen4_exp/runtime.py b/unit_tests/models/qwen4_exp/runtime.py new file mode 100644 index 0000000000..0a02d45d4e --- /dev/null +++ b/unit_tests/models/qwen4_exp/runtime.py @@ -0,0 +1,133 @@ +import json +from types import SimpleNamespace +from contextlib import nullcontext +from uuid import uuid4 + +import pytest +import torch + +from lightllm.models.qwen4_exp.layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight +from lightllm.models.qwen4_exp.infer_struct import Qwen4ExpInferStateInfo +from lightllm.models.qwen4_exp.model import Qwen4ExpTpPartModel +from lightllm.common.quantization import Quantcfg +from lightllm.utils import envs_utils +from .helpers import checkpoint, prepost_checkpoint + + +@pytest.fixture +def runtime(config, tmp_path, monkeypatch, request): + (tmp_path / "config.json").write_text(json.dumps(config)) + args = envs_utils.get_env_start_args() + args.model_dir = str(tmp_path) + args.data_type = "bf16" + for name, value in getattr(request, "param", {}).items(): + setattr(args, name, value) + args.cpu_cache_token_page_size = args.linear_att_hash_page_size * args.linear_att_page_block_num + args.linear_att_ssm_data_type = "float32" + args.llm_prefill_att_backend = ["triton", "triton"] + args.llm_decode_att_backend = ["triton", "triton"] + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(vars(args))) + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") + monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", "qwen-runtime-test-" + uuid4().hex) + envs_utils.get_env_start_args.cache_clear() + envs_utils.get_unique_server_name.cache_clear() + weight = Qwen4ExpTransformerLayerWeight(1, torch.bfloat16, config, Quantcfg(config)) + model = Qwen4ExpTpPartModel.__new__(Qwen4ExpTpPartModel) + model.__dict__.update( + config=config, + args=args, + tp_world_size_=1, + tp_rank_=0, + torch_memory_saver=SimpleNamespace(region=lambda **kwargs: nullcontext()), + data_type=torch.bfloat16, + max_req_num=4, + max_seq_length=64, + batch_max_tokens=8, + max_total_token_num=32, + num_kv_heads=2, + mem_fraction=0.9, + trans_layers_weight=[weight], + ) + from lightllm.common.req_manager import ReqManager + from .resource_helpers import bind_states + + model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(config["n_layer"])] + req = ReqManager(4, 64) + model.bind_req_manager(req) + model.prepare_attention_states() + mem = bind_states(model) + assert not weight.ngram_hasher.verify_load() + torch.manual_seed(1307) + source = checkpoint(config, 1) + weight.load_hf_weights(dict(source)) + weight.verify_load() + angle = torch.randn(64, 2, device="cuda") + + model._cos_cached = angle.cos().bfloat16() + model._sin_cached = angle.sin().bfloat16() + model.mtp_manager = SimpleNamespace(get_decode_draft_step=lambda is_draft: 0) + model._init_att_backend() + model._init_att_backend1() + yield SimpleNamespace(mem=mem, req=req, weight=weight, source=source, model=model) + envs_utils.get_unique_server_name.cache_clear() + + +def state_for(runtime, requests, lengths, tokens, ready=None, mtp_indices=None): + state = Qwen4ExpInferStateInfo() + state.req_manager = runtime.req + state.att_state_manager = runtime.model.att_state_manager + state.att_state_role = runtime.model.att_state_role + state.is_prefill = ready is not None + state.batch_size = len(requests) + state.b_req_idx = torch.tensor(requests, device="cuda", dtype=torch.int32) + state.b_seq_len = torch.tensor(lengths, device="cuda", dtype=torch.int32) + if mtp_indices is not None: + state.b_mtp_index = torch.tensor(mtp_indices, device="cuda", dtype=torch.int32) + state.input_ids = torch.tensor(tokens, device="cuda", dtype=torch.int64) + state.mem_index = runtime.mem.alloc(len(tokens)).cuda(non_blocking=True) + if ready is not None: + state.b_ready_cache_len = torch.tensor(ready, device="cuda", dtype=torch.int32) + query_lengths = [length - prefix for length, prefix in zip(lengths, ready)] + else: + state.b_position_delta = torch.zeros(len(requests), device="cuda", dtype=torch.int64) + query_lengths = [1] * len(requests) + offset = 0 + for request, length, query in zip(requests, lengths, query_lengths): + runtime.req.req_to_token_indexs[request, length - query : length] = state.mem_index[offset : offset + query] + offset += query + state.multimodal_params = [] + state.init_some_extra_state(runtime.model) + if state.is_prefill: + state.prefill_att_state = runtime.model.prefill_att_backend.create_att_prefill_state(state) + if runtime.model.prefill_att_backend1 is not None: + state.prefill_att_state1 = runtime.model.prefill_att_backend1.create_att_prefill_state(state) + else: + state.decode_att_state = runtime.model.decode_att_backend.create_att_decode_state(state) + if runtime.model.decode_att_backend1 is not None: + state.decode_att_state1 = runtime.model.decode_att_backend1.create_att_decode_state(state) + state.init_att_state() + return state + + +def assemble_model(runtime, config): + model, source = runtime.model, prepost_checkpoint(config) + weights = [] + for layer in range(config["n_layer"]): + if layer == 1: + weight = runtime.weight + source.update(runtime.source) + else: + local = checkpoint(config, layer) + source.update(local) + weight = model.transformer_weight_class(layer, torch.bfloat16, config, Quantcfg(config)) + weight.load_hf_weights(dict(local)) + weight.verify_load() + weights.append(weight) + model.trans_layers_weight = weights + model.pre_post_weight = model.pre_and_post_weight_class(torch.bfloat16, config) + model.pre_post_weight.load_hf_weights(dict(source)) + model.pre_post_weight.verify_load() + model._init_infer_layer() + model._init_some_value() + model.prefill_graph = None + return model, source diff --git a/unit_tests/models/qwen4_exp/test_attention_state.py b/unit_tests/models/qwen4_exp/test_attention_state.py new file mode 100644 index 0000000000..62cce52c92 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_attention_state.py @@ -0,0 +1,205 @@ +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state.manager import AttentionStateManager +from lightllm.common.attention_state.types import StateTransferLayout +from lightllm.models import get_model_class +from lightllm.models.draft_registry import get_draft_model_class +from lightllm.models.qwen4_exp.model import Qwen4ExpTpPartModel +from lightllm.models.qwen4_exp.mtp.model import Qwen4ExpMTPModel +from lightllm.models.qwen4_exp.attention_state import Qwen4ExpAttentionState +from lightllm.models.qwen4_exp.lexical_state import Qwen4ExpLexicalState, LexicalStateConfig +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.common.attention_state.bounded_att.gdn.gdn_state import GDNState +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState + + +@pytest.fixture +def lexical(monkeypatch): + import lightllm.models.qwen4_exp.lexical_state as module + + def copy(*, copies): + for source, destination in copies: + destination.copy_(source) + + monkeypatch.setattr(module, "copy_bounded_att_tensors", copy) + state = Qwen4ExpLexicalState(LexicalStateConfig((1, 5), 2, 6, 8, 63, torch.bfloat16), 4, 2, "cpu") + state.state_name = "main_qwen4_lexical_state" + state.initialize_page_pools(4, 2, pin_memory=False) + return state + + +def test_lazy_registration_and_resource_factories(config): + args = get_env_start_args() + args.data_type, args.linear_att_ssm_data_type = "bf16", "float32" + args.cpu_cache_token_page_size = 8 + args.pd_kv_page_size = 4 + for model_type in ("qwen4_exp", "qwen4_exp_text"): + assert get_model_class({"model_type": model_type}) is Qwen4ExpTpPartModel + assert get_draft_model_class({"model_type": model_type}, "eagle_with_att") is Qwen4ExpMTPModel + manager = AttentionStateManager(request_capacity=4) + for cls, role, layers in ((Qwen4ExpTpPartModel, "main", (0, 1, 2, 3)), (Qwen4ExpMTPModel, "draft0", (4,))): + full = cls.create_full_att_states(config, args, layers) + bounded = cls.create_bounded_att_states(config, args, layers, 4) + for state in (*full, *bounded): + manager.register_state(role, state) + manager.finish_state_registration() + state = manager.get_state("main", "", 3, Qwen4ExpAttentionState) + assert isinstance(state, Qwen4ExpAttentionState) + assert manager.get_state("draft0", "qwen4_gqa", 4) is not state + assert len(manager.get_bounded_attention_states("main")) == 2 + assert not manager.get_bounded_attention_states("draft0") + for state in manager.get_full_attention_states(): + regions = state.get_cpu_transfer_regions() + assert len(regions) == 6 + assert all(left.byte_end <= right.byte_start for left, right in zip(regions, regions[1:])) + assert {region.field for region in regions} == { + "k", + "v", + "indexer.k", + "indexer.v", + "positions.k", + "positions.v", + } + assert state.get_token_bytes() > 0 + + +@pytest.mark.parametrize("mode,state_type", [("native", GDNState), ("replay", ReplayGDNState)]) +def test_mtp_state_modes_keep_qwen4_resources(config, mode, state_type): + args = get_env_start_args() + args.data_type = "bfloat16" + args.mtp_step = 3 + args.ssm_state_mode = mode + args.ssm_replay_config = {"cache_len": 16, "compute_mode": "precompute"} if mode == "replay" else {} + args.linear_att_ssm_data_type = "float32" + layers = (0, 1, 2, 3) + bounded = Qwen4ExpTpPartModel.create_bounded_att_states(config, args, layers, 4, device="cpu") + manager = AttentionStateManager(request_capacity=4) + for state in (*Qwen4ExpTpPartModel.create_full_att_states(config, args, layers), *bounded): + manager.register_state("main", state) + manager.finish_state_registration() + state = manager.get_state("main", "gdn", 1, state_type) + assert state.config.logical_layer_ids == (0, 1, 2) + assert state.get_layer_view(1)[1].dtype == torch.float32 + assert isinstance(manager.get_state("main", "qwen4_gqa", 3), Qwen4ExpAttentionState) + assert any(isinstance(state, Qwen4ExpLexicalState) for state in bounded) + assert not Qwen4ExpMTPModel.create_bounded_att_states(config, args, (4,), 4, device="cpu") + + +@pytest.mark.parametrize("tp", [1, 2, 4, 8]) +def test_replicated_indexer_and_positions_layout(config, tp): + args = get_env_start_args() + args.data_type = "bf16" + for rank in range(tp): + (state,) = Qwen4ExpTpPartModel.create_full_att_states(config, args, (0, 1, 2, 3), rank, tp) + regions = state._get_transfer_regions(8) + for region in regions[2:]: + assert region.write_enabled == (rank == 0) + assert region.shape[2] == 1 + assert sum(region.byte_end - region.byte_start for region in regions) == regions[-1].byte_end + + +def test_state_storage_accounts_for_all_fields_and_restores_positions(config): + args = get_env_start_args() + args.data_type = "bf16" + (state,) = Qwen4ExpTpPartModel.create_full_att_states(config, args, (0, 1, 2, 3)) + state.device = torch.device("cpu") + state.allocate_storage(8) + tensors = ( + state.k_storage, + state.v_storage, + state.indexer.k_storage, + state.indexer.v_storage, + state.positions.k_storage, + state.positions.v_storage, + ) + assert sum(t.numel() * t.element_size() for t in tensors) == (8 + state.page_size) * state.get_token_bytes() + indices = torch.tensor([1, 4]) + cos = torch.arange(8).view(2, 4).bfloat16() + state.store_positions(indices, cos, -cos) + view = state.qsa_view(3) + assert torch.equal(view.key_cos[indices], cos) + assert torch.equal(view.key_sin[indices], -cos) + assert view.raw_keys.data_ptr() == state.indexer.k_storage.data_ptr() + + +@pytest.mark.parametrize("page_kind", ["big", "small"]) +def test_lexical_prefix_round_trip_is_independent_of_runtime(lexical, page_kind): + lexical.ngram[:, 1].fill_(17) + lexical.convolution[:, 1].fill_(3) + lexical.save_bounded_att_runtime_to_pages([1], [128], page_kind, [0]) + lexical.reset_bounded_att_state([1]) + assert bool((lexical.ngram[:, 1] == 63).all()) + lexical.restore_bounded_att_runtime_from_pages([2], [128], page_kind, [0]) + assert bool((lexical.ngram[:, 2] == 17).all()) + assert bool((lexical.convolution[:, 2] == 3).all()) + lexical.reset_bounded_att_state([2]) + assert bool((lexical._pool(page_kind).ngram[0] == 17).all()) + + +def layout_for(state): + regions = state.get_cpu_transfer_regions() + return StateTransferLayout( + state.state_name, + None, + 16, + 16 + regions[-1].byte_end, + tuple(replace(region, byte_start=region.byte_start + 16, byte_end=region.byte_end + 16) for region in regions), + ) + + +def test_cpu_and_pd_lexical_transfers_preserve_unrelated_bytes(lexical): + layout = layout_for(lexical) + pages = torch.full((3, layout.byte_end + 16), 0x5A, dtype=torch.uint8) + lexical.ngram[:, 1].fill_(7) + lexical.convolution[:, 1].fill_(11) + lexical.save_bounded_att_runtime_to_pages([1], [7], "big", [0]) + lexical.write_state_pages_to_cpu("big", [0, 0], pages, [1, 2], layout, [False, True]) + assert bool((pages[2] == 0x5A).all()) + assert bool((pages[:, :16] == 0x5A).all()) + assert bool((pages[:, layout.byte_end :] == 0x5A).all()) + lexical.read_state_pages_from_cpu(pages, [1], [2], layout) + lexical.restore_bounded_att_runtime_from_pages([2], [7], "big", [2]) + assert torch.equal(lexical.ngram[:, 2], lexical.ngram[:, 1]) + assert torch.equal(lexical.convolution[:, 2], lexical.convolution[:, 1]) + lexical.write_to_pd_page(1, pages, 0, 7, layout) + lexical.reset_bounded_att_state([2]) + lexical.read_from_pd_page(pages, 2, 0, 7, layout) + assert torch.equal(lexical.ngram[:, 2], lexical.ngram[:, 1]) + assert torch.equal(lexical.convolution[:, 2], lexical.convolution[:, 1]) + + +def test_dp_copy_includes_runtime_and_prefix(lexical): + other = Qwen4ExpLexicalState(lexical.config, 4, 2, "cpu") + other.initialize_page_pools(4, pin_memory=False) + other.ngram[:, 0].fill_(29) + other.convolution[:, 0].fill_(13) + lexical.copy_runtime_state_from_dp([0], [2], other) + assert bool((lexical.ngram[:, 2] == 29).all()) + other.save_bounded_att_runtime_to_pages([0], [8], "big", [1]) + lexical.copy_big_pages_from_dp([1], [3], other) + lexical.restore_bounded_att_runtime_from_pages([1], [8], "big", [3]) + assert bool((lexical.convolution[:, 1] == 13).all()) + + +def test_mtp_commit_maps_request_rows_and_keeps_hold(lexical): + observed = [] + lexical.speculative_lexical_caches = {1: SimpleNamespace(commit=lambda *args: observed.append(args))} + requests, accepted = torch.tensor([2, 0, 3]), torch.tensor([1, 3, 1]) + lexical.commit_bounded_att_state(requests, accepted) + req, starts, lengths, hold = observed[0] + assert req is requests and lengths is accepted + assert starts.tolist() == [0, 1, 2] and hold == 3 + + +def test_ipc_state_excludes_weights_and_speculative_workspaces(lexical): + lexical.lexical_caches = {1: SimpleNamespace(layer_weight=object())} + lexical.speculative_lexical_caches = {1: object()} + exported = lexical.__getstate__() + assert not exported["lexical_caches"] and not exported["speculative_lexical_caches"] + assert exported["ngram"] is lexical.ngram + assert exported["convolution"] is lexical.convolution + assert lexical.lexical_caches and lexical.speculative_lexical_caches diff --git a/unit_tests/models/qwen4_exp/test_batch.py b/unit_tests/models/qwen4_exp/test_batch.py new file mode 100644 index 0000000000..e4895c294d --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_batch.py @@ -0,0 +1,97 @@ +import pytest +import torch + +from lightllm.models.qwen4_exp.batch import PackedBatchMetadata +from lightllm.models.qwen4_exp.state import QSABatch +from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def prepare(requests, lengths, query_lengths, hold=7): + tensors = [torch.tensor(value, device="cuda", dtype=torch.int32) for value in (requests, lengths, query_lengths)] + return tensors, PackedBatchMetadata.prepare(*tensors, sum(query_lengths), hold, 4) + + +def check(metadata, requests, lengths, query_lengths, hold=7): + counts = list(query_lengths) + for row, request in enumerate(requests): + if request != hold and row + 1 < len(requests) and requests[row + 1] == request: + counts[row + 1] += counts[row] + counts[row] = 0 + starts, active, key_lengths, first, rows, visible, mask = [0], [], [], [], [], [], [] + for request, length, count in zip(requests, lengths, counts): + real = request != hold and count > 0 + starts.append(starts[-1] + count) + active.append(request if real else -1) + key_lengths.append(length if real else 0) + first.append((length - count) // 4 if real else 0) + rows.extend([request] * count) + visible.extend(range(length - count + 1, length + 1) if real else [0] * count) + mask.extend([real] * count) + expected = [starts, active, key_lengths, first, rows, visible, mask] + for name, values in zip(metadata.__dataclass_fields__, expected): + actual = getattr(metadata, name) + torch.testing.assert_close(actual, torch.tensor(values, device="cuda", dtype=actual.dtype)) + + +@pytest.mark.parametrize( + "requests,lengths,query_lengths", + [ + ([5, 2, 7], [13, 21, 4], [13, 5, 4]), + ([5, 2, 7, 7], [14, 22, 2, 2], [1, 1, 1, 1]), + ([5, 5, 5, 2, 2, 7, 7], [14, 15, 16, 22, 23, 2, 2], [1] * 7), + ([5, 2, 7], [13, 21, 0], [0, 0, 0]), + ([5, 2], [1500, 1900], [1200, 1700]), + ([], [], []), + ], +) +def test_native_packed_batch_metadata(requests, lengths, query_lengths): + _, metadata = prepare(requests, lengths, query_lengths) + check(metadata, requests, lengths, query_lengths) + tokens = sum(query_lengths) + slots = torch.zeros(8, 2000, device="cuda", dtype=torch.int32) + memory = torch.arange(tokens, device="cuda", dtype=torch.int32) + rotary = torch.ones(tokens, 4, device="cuda", dtype=torch.bfloat16) + batch = QSABatch.from_metadata(metadata, slots, memory, rotary, rotary, 4) + assert batch.slot_table.data_ptr() == slots.data_ptr() + assert batch.active_rows.data_ptr() == metadata.active_rows.data_ptr() + assert batch.mem_indices.dtype == torch.int64 + assert batch.max_new_blocks == (tokens + 3) // 4 + + +def test_metadata_visibility_matches_existing_native_speculative_layout(): + requests, lengths, query_lengths = [5, 5, 5, 2, 2, 7, 7], [14, 15, 16, 22, 23, 2, 2], [1] * 7 + tensors, metadata = prepare(requests, lengths, query_lengths) + slots = torch.arange(8 * 32, device="cuda", dtype=torch.int32).reshape(8, 32) + ragged = torch.empty(sum([16, 23, 2, 2]), device="cuda", dtype=torch.int32) + _, _, native_lengths = gen_nsa_ks_ke(tensors[1], tensors[2], tensors[0], slots, 7, ragged, hold_req_idx=7) + torch.testing.assert_close(metadata.visible_lengths[metadata.token_mask], native_lengths[metadata.token_mask]) + + +def test_metadata_graph_replay_and_copy_preserve_buffer_addresses(): + requests, lengths, query_lengths = [5, 5, 2, 7, 7], [14, 15, 22, 2, 2], [1] * 5 + tensors, metadata = prepare(requests, lengths, query_lengths) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = PackedBatchMetadata.prepare(*tensors, 5, 7, 4) + pointers = {name: getattr(captured, name).data_ptr() for name in captured.__dataclass_fields__} + for ids, sizes in [([1, 3, 5, 7, 7], [25, 41, 13, 2, 2]), ([2, 2, 2, 5, 7], [12, 13, 14, 5, 2])]: + tensors[0].copy_(torch.tensor(ids, device="cuda", dtype=torch.int32)) + tensors[1].copy_(torch.tensor(sizes, device="cuda", dtype=torch.int32)) + graph.replay() + check(captured, ids, sizes, query_lengths) + metadata.copy_from(captured) + check(metadata, ids, sizes, query_lengths) + assert pointers == {name: getattr(captured, name).data_ptr() for name in captured.__dataclass_fields__} + _, incompatible = prepare([1], [1], [1]) + with pytest.raises(ValueError, match="Incompatible graph batch metadata"): + metadata.copy_from(incompatible) + + +def test_metadata_rejects_incompatible_input_shapes(): + values = torch.ones(2, device="cuda", dtype=torch.int32) + with pytest.raises(ValueError, match="matching contiguous"): + PackedBatchMetadata.prepare(values, values[:1], values, 2, 7, 4) + with pytest.raises(ValueError, match="compression ratio"): + PackedBatchMetadata.prepare(values, values, values, 2, 7, 0) diff --git a/unit_tests/models/qwen4_exp/test_decoder.py b/unit_tests/models/qwen4_exp/test_decoder.py new file mode 100644 index 0000000000..fc466ec680 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_decoder.py @@ -0,0 +1,214 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.qwen4_exp.layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight +from lightllm.models.qwen4_exp.layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer +from lightllm.models.qwen4_exp.state import QSABatch, QSAExecutionState +from lightllm.models.qwen4_exp.batch import PackedBatchMetadata +from lightllm.models.qwen4_exp.position import combine_mrope +from lightllm.common.quantization import Quantcfg +from lightllm.models.qwen4_exp.attention_state import Qwen4ExpAttentionState +from lightllm.common.attention_state.full_att.gqa.gqa_state import GQAAttentionConfig +from .helpers import checkpoint +from .reference import reference_layer + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def setup(config, lengths, physical_cache=False): + source = checkpoint(config, 3) + weight = Qwen4ExpTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + weight.load_hf_weights(dict(source)) + weight.verify_load() + infer = Qwen4ExpTransformerLayerInfer(3, config) + + from lightllm.utils.dist_utils import get_dp_world_size + + manager = Qwen4ExpAttentionState( + GQAAttentionConfig( + logical_layer_groups=((3,),), + global_head_num=config["num_key_value_heads"], + head_dim=config["head_dim"], + dtype=torch.bfloat16, + tp_world_size=get_dp_world_size(), + ), + 16, + 4, + device="cuda", + ) + manager.allocate_storage(47) + rows = [5, 2] + table = torch.full((8, 20), -1, device="cuda", dtype=torch.int32) + slots = torch.randperm(48, device="cuda")[: sum(lengths)] + offset = 0 + for row, count in zip(rows, lengths): + table[row, :count] = slots[offset : offset + count].int() + offset += count + dtype = torch.bfloat16 + angle = torch.randn(3, sum(lengths), 2, device="cuda") + pc, ps = angle.cos().to(dtype), angle.sin().to(dtype) + cos, sin = combine_mrope(pc, ps, [1, 1, 0]) + key_cos, key_sin = torch.zeros(48, 4, device="cuda", dtype=dtype), torch.zeros(48, 4, device="cuda", dtype=dtype) + key_cos.index_copy_(0, slots, cos) + key_sin.index_copy_(0, slots, sin) + metadata = PackedBatchMetadata.prepare( + torch.tensor([5, 7, 2], device="cuda", dtype=torch.int32), + torch.tensor([lengths[0], 0, lengths[1]], device="cuda", dtype=torch.int32), + torch.tensor([lengths[0], 0, lengths[1]], device="cuda", dtype=torch.int32), + sum(lengths), + 7, + 4, + ) + batch = QSABatch.from_metadata(metadata, table, slots, cos, sin, 4) + view = QSAExecutionState( + torch.zeros(48, 16, device="cuda", dtype=dtype), + torch.full((48, 16) if physical_cache else (8, 5, 16), float("nan"), device="cuda", dtype=dtype), + key_cos, + key_sin, + ) + state = SimpleNamespace( + is_prefill=True, + need_dp_prefill_balance=False, + dist_group=None, + position_cos=pc, + position_sin=ps, + mem_index=slots.int(), + att_state_role="main", + att_state_manager=SimpleNamespace(get_state=lambda *args, **kwargs: manager), + qwen4_exp_qsa={3: view}, + qwen4_exp_batch=batch, + ) + return infer, weight, source, state + + +@pytest.mark.parametrize("physical_cache", [False, True]) +def test_native_qsa_decoder_matches_reference(config, physical_cache): + torch.manual_seed(937) + lengths = [13, 5] + infer, weights, source, state = setup(config, lengths, physical_cache) + residual = torch.randn(sum(lengths), 64, device="cuda", dtype=torch.bfloat16) + before = residual.clone() + actual = infer.context_forward(residual, state, weights) + expected = reference_layer(residual, source, config, state.qwen4_exp_batch.cos, state.qwen4_exp_batch.sin, lengths) + torch.testing.assert_close(actual, expected, atol=0.016, rtol=0.016) + torch.testing.assert_close(residual, before, atol=0, rtol=0) + key, value = state.att_state_manager.get_state().get_layer_view(3) + assert key.shape == value.shape + + old_batch = state.qwen4_exp_batch + + used = set(old_batch.mem_indices.tolist()) + slots = torch.tensor([slot for slot in range(48) if slot not in used][:2], device="cuda", dtype=torch.int64) + angle = torch.randn(3, 2, 2, device="cuda") + pc, ps = angle.cos().to(residual.dtype), angle.sin().to(residual.dtype) + cos, sin = combine_mrope(pc, ps, [1, 1, 0]) + view = state.qwen4_exp_qsa[3] + view.key_cos.index_copy_(0, slots, cos) + view.key_sin.index_copy_(0, slots, sin) + for row, count, slot in zip([5, 2], lengths, slots): + old_batch.slot_table[row, count] = slot.int() + metadata = PackedBatchMetadata.prepare( + torch.tensor([5, 2], device="cuda", dtype=torch.int32), + torch.tensor([14, 6], device="cuda", dtype=torch.int32), + torch.ones(2, device="cuda", dtype=torch.int32), + 2, + 7, + 4, + ) + state.qwen4_exp_batch = QSABatch.from_metadata(metadata, old_batch.slot_table, slots, cos, sin, 4) + state.position_cos, state.position_sin, state.mem_index, state.is_prefill = pc, ps, slots.int(), False + new_residual = torch.randn(2, 64, device="cuda", dtype=residual.dtype) + actual = infer.token_forward(new_residual, state, weights) + all_residual = torch.cat([residual[:13], new_residual[:1], residual[13:], new_residual[1:]]) + all_cos = torch.cat([old_batch.cos[:13], cos[:1], old_batch.cos[13:], cos[1:]]) + all_sin = torch.cat([old_batch.sin[:13], sin[:1], old_batch.sin[13:], sin[1:]]) + full_reference = reference_layer(all_residual, source, config, all_cos, all_sin, [14, 6]) + torch.testing.assert_close(actual, full_reference[[13, 19]], atol=0.016, rtol=0.016) + + +@pytest.mark.parametrize("physical_cache", [False, True]) +def test_native_qsa_decoder_cuda_graph_replay(config, physical_cache): + torch.manual_seed(937) + lengths = [13, 5] + infer, weights, source, state = setup(config, lengths, physical_cache) + residual = torch.randn(sum(lengths), 64, device="cuda", dtype=torch.bfloat16) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + infer.context_forward(residual, state, weights) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = infer.context_forward(residual, state, weights) + batch = state.qwen4_exp_batch + view = state.qwen4_exp_qsa[3] + for _ in range(2): + residual.normal_() + angle = torch.randn_like(state.position_cos.float()) + state.position_cos.copy_(angle.cos()) + state.position_sin.copy_(angle.sin()) + cos, sin = combine_mrope(state.position_cos, state.position_sin, [1, 1, 0]) + batch.cos.copy_(cos) + batch.sin.copy_(sin) + view.key_cos.index_copy_(0, batch.mem_indices, cos) + view.key_sin.index_copy_(0, batch.mem_indices, sin) + graph.replay() + expected = reference_layer(residual, source, config, cos, sin, lengths) + torch.testing.assert_close(actual, expected, atol=0.016, rtol=0.016) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("prefill", [False, True]) +def test_long_decode_attention_matches_single_queries(dtype, prefill): + from lightllm.common.qsa.indexer import QSAIndexer + + torch.manual_seed(701) + slots, rows, dim = 4096, 3, 128 + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=dtype) + + indexer = QSAIndexer(rand(5 * dim, dim), rand(dim), rand(dim)) + raw = rand(slots, dim) + pooled = rand(1, slots // 4, dim) + angles = torch.empty(slots, 0, device="cuda", dtype=dtype) + view = QSAExecutionState(raw, pooled, angles, angles) + query, key, value = rand(rows, 6, 256), rand(slots, 1, 256), rand(slots, 1, 256) + hidden = rand(rows, dim) + table = torch.arange(slots, device="cuda", dtype=torch.int32)[None, :] + + def execute(start, end, decode=True, attention_splits=None): + count = end - start + batch = QSABatch( + table, + torch.tensor([0], device="cuda", dtype=torch.int32), + torch.tensor([slots], device="cuda", dtype=torch.int32), + torch.tensor([slots // 4], device="cuda", dtype=torch.int32), + torch.zeros(count, device="cuda", dtype=torch.int32), + torch.arange(2050 + start, 2050 + end, device="cuda", dtype=torch.int32), + torch.arange(slots - rows + start, slots - rows + end, device="cuda"), + angles[:count], + angles[:count], + 0, + ) + return view.attend( + indexer, + hidden[start:end], + query[start:end], + key, + value, + batch, + decode=decode, + attention_splits=attention_splits, + ) + + options = {"decode": not prefill, "attention_splits": None} + expected = torch.cat([execute(row, row + 1, **options) for row in range(rows)]) + if prefill: + torch.testing.assert_close(execute(0, rows, **options), expected, atol=1e-5, rtol=1e-5) + else: + torch.testing.assert_close(execute(0, rows, **options), expected, atol=0, rtol=0) + torch.testing.assert_close(execute(0, 1, decode=False), expected[:1], atol=0, rtol=0) diff --git a/unit_tests/models/qwen4_exp/test_embedding_residency.py b/unit_tests/models/qwen4_exp/test_embedding_residency.py new file mode 100644 index 0000000000..9ca54c0478 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_embedding_residency.py @@ -0,0 +1,45 @@ +import argparse + +import pytest +import torch + +from lightllm.models.qwen4_exp.layer_weights import embedding_weight as module +from lightllm.server.api_cli import add_cli_args +from lightllm.server.core.objs.start_args_type import StartArgs + + +def test_residency_is_opt_in_and_cli_matches_start_args(): + parser = add_cli_args(argparse.ArgumentParser()) + assert not parser.parse_args([]).offload_ngram_embedding + assert not StartArgs().offload_ngram_embedding + assert parser.parse_args(["--offload_ngram_embedding"]).offload_ngram_embedding + + +@pytest.mark.parametrize("args", [{}, {"offload_ngram_embedding": False}]) +def test_gpu_factory_preserves_existing_constructor(args, monkeypatch): + monkeypatch.setattr(module, "get_env_start_args", lambda: args) + monkeypatch.setattr(module, "RowShardedEmbeddingWeight", lambda *values: values) + result = module.ngram_embedding_weight(160, 1024, "table.weight", torch.bfloat16, 4) + assert result == (160, 1024, "table.weight", torch.bfloat16, 4) + + +@pytest.mark.parametrize("torch_fallback", [False, True]) +def test_cpu_factory_injects_model_owned_allocation_policy(monkeypatch, torch_fallback): + monkeypatch.setattr( + module, + "get_env_start_args", + lambda: {"offload_ngram_embedding": True, "hardware_platform": "cuda", "enable_torch_fallback": torch_fallback}, + ) + monkeypatch.setattr(module, "PinnedRowShardedEmbeddingWeight", lambda *values, **kwargs: (values, kwargs)) + from lightllm.models.qwen4_exp.pinned_memory import pinned_zeros + + values, kwargs = module.ngram_embedding_weight(160, 1024, "table.weight", torch.bfloat16, 4) + assert values == (160, 1024, "table.weight", torch.bfloat16, 4) + assert kwargs == {"cpu_allocator": pinned_zeros} + + +@pytest.mark.parametrize("extra", [{"hardware_platform": "musa"}, {"hardware_platform": "cpu"}]) +def test_unsupported_offload_configuration_fails_before_allocation(extra, monkeypatch): + monkeypatch.setattr(module, "get_env_start_args", lambda: dict(offload_ngram_embedding=True, **extra)) + with pytest.raises(ValueError, match="CUDA Triton"): + module.ngram_embedding_weight(160, 1024, "table.weight", torch.bfloat16, 4) diff --git a/unit_tests/models/qwen4_exp/test_eos.py b/unit_tests/models/qwen4_exp/test_eos.py new file mode 100644 index 0000000000..31e15627d9 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_eos.py @@ -0,0 +1,20 @@ +import json +from types import SimpleNamespace + +import pytest +from transformers import AutoTokenizer + +from lightllm.utils.config_utils import get_eos_token_ids + + +@pytest.mark.parametrize("model_type", ["qwen4_exp", "qwen4_exp_text", "qwen3_5", "llama"]) +@pytest.mark.parametrize("generation_eos", [[248046, 248044], 248046, None]) +def test_chat_end_and_end_of_text_are_preserved(tmp_path, monkeypatch, model_type, generation_eos): + (tmp_path / "config.json").write_text( + json.dumps({"model_type": model_type, "text_config": {"eos_token_id": 248044}}) + ) + if generation_eos is not None: + (tmp_path / "generation_config.json").write_text(json.dumps({"eos_token_id": generation_eos})) + monkeypatch.setattr(AutoTokenizer, "from_pretrained", lambda *args, **kwargs: SimpleNamespace(eos_token_id=248046)) + expected = {248044} if model_type == "llama" else {248044, 248046} + assert set(get_eos_token_ids(str(tmp_path))) == expected diff --git a/unit_tests/models/qwen4_exp/test_lexical_cache.py b/unit_tests/models/qwen4_exp/test_lexical_cache.py new file mode 100644 index 0000000000..fef7202721 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_lexical_cache.py @@ -0,0 +1,100 @@ +import pytest +import torch + +from lightllm.models.qwen4_exp.lexical_cache import LexicalSnapshot +from .lexical_helpers import make_cache, metadata, reference_segment + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("capture", [False, True]) +def test_request_histories_acceptance_padding_and_graphs(config, capture): + torch.manual_seed(1207) + cache, source = make_cache(config) + + cache.ngram[5].copy_(torch.tensor([7, 11], device="cuda")) + cache.convolution[5].normal_() + prefix = cache.snapshot(5) + cache.restore(2, prefix) + requests = [5, 5, 5, 2, 7, 7] + batch_meta = metadata(requests, [14, 15, 16, 14, 2, 2], [1] * 6) + accepted = torch.tensor([0, 0, 2, 1, 0, 0], device="cuda", dtype=torch.int32) + tokens = torch.zeros(6, device="cuda", dtype=torch.int64) + residual = torch.randn(6, 64, device="cuda", dtype=torch.bfloat16) + untouched_ngram, untouched_conv = cache.ngram[[0, 7]].clone(), cache.convolution[[0, 7]].clone() + + def run(): + candidate = cache.prepare(batch_meta) + update = candidate.apply(residual, tokens, cache.layer_weight, 1, None) + cache.commit(batch_meta, candidate, accepted) + return update + + if capture: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + result = run() + cache.restore(5, prefix) + cache.restore(2, prefix) + ref_ngram = {request: prefix.ngram[None].clone() for request in (5, 2)} + ref_conv = {request: prefix.convolution[None].clone() for request in (5, 2)} + for ids in ([19, 63, 23, 29, 31, 37], [41, 43, 47, 63, 5, 7], [11, 13, 17, 19, 23, 29]): + tokens.copy_(torch.tensor(ids, device="cuda", dtype=torch.int64)) + residual.normal_() + if capture: + graph.replay() + actual = result + else: + before_ngram, before_conv = cache.ngram.clone(), cache.convolution.clone() + candidate = cache.prepare(batch_meta) + actual = candidate.apply(residual, tokens, cache.layer_weight, 1, None) + + torch.testing.assert_close(cache.ngram, before_ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution, before_conv, atol=0, rtol=0) + cache.commit(batch_meta, candidate, accepted) + for request, start, end, count in ((5, 0, 3, 2), (2, 3, 4, 1)): + expected, states = reference_segment( + residual[start:end], tokens[start:end], source, config, ref_ngram[request], ref_conv[request] + ) + torch.testing.assert_close(actual[start:end], expected, atol=0.016, rtol=0.016) + ref_ngram[request], ref_conv[request] = states[count] + torch.testing.assert_close(cache.ngram[request], ref_ngram[request][0], atol=0, rtol=0) + torch.testing.assert_close(cache.convolution[request], ref_conv[request][0], atol=0.016, rtol=0.016) + torch.testing.assert_close(cache.ngram[[0, 7]], untouched_ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution[[0, 7]], untouched_conv, atol=0, rtol=0) + torch.testing.assert_close(prefix.ngram, torch.tensor([7, 11], device="cuda"), atol=0, rtol=0) + + +def test_lexical_snapshot_reset_cpu_restore_and_rejection(config): + torch.manual_seed(1207) + cache, _ = make_cache(config) + cache.ngram[3].copy_(torch.tensor([7, 11], device="cuda")) + cache.convolution[3].normal_() + snapshot = cache.snapshot(3) + cache.reset(3) + assert torch.all(cache.ngram[3] == config["eos_token_id"]) + assert torch.count_nonzero(cache.convolution[3]) == 0 + cache.restore(5, LexicalSnapshot(snapshot.ngram.cpu(), snapshot.convolution.cpu())) + torch.testing.assert_close(cache.ngram[5], snapshot.ngram) + torch.testing.assert_close(cache.convolution[5], snapshot.convolution) + before_ngram, before_conv = cache.ngram.clone(), cache.convolution.clone() + batch_meta = metadata([5], [12], [3]) + candidate = cache.prepare(batch_meta) + candidate.apply( + torch.randn(3, 64, device="cuda", dtype=torch.bfloat16), + torch.tensor([17, 63, 19], device="cuda"), + cache.layer_weight, + 1, + None, + ) + cache.commit(batch_meta, candidate, torch.zeros(1, device="cuda", dtype=torch.int32)) + torch.testing.assert_close(cache.ngram, before_ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution, before_conv, atol=0, rtol=0) + with pytest.raises(ValueError, match="incompatible"): + cache.restore(3, LexicalSnapshot(snapshot.ngram[:1], snapshot.convolution)) + assert torch.all(cache.ngram[3] == config["eos_token_id"]) diff --git a/unit_tests/models/qwen4_exp/test_linear_decoder.py b/unit_tests/models/qwen4_exp/test_linear_decoder.py new file mode 100644 index 0000000000..952c268a3a --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_linear_decoder.py @@ -0,0 +1,183 @@ +from argparse import Namespace +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.qwen4_exp.layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight +from lightllm.models.qwen4_exp.layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer +from lightllm.models.qwen4_exp.state import PLEBatch +from lightllm.models.qwen4_exp.batch import PackedBatchMetadata +from lightllm.models.qwen4_exp.lexical_cache import LexicalRequestCache +from lightllm.common.quantization import Quantcfg +from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend +from .helpers import checkpoint +from .reference import reference_step + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("capture", [False, True]) +def test_native_ple_gdn_and_moe_cached_steps(config, capture): + torch.manual_seed(1037) + source = checkpoint(config, 1) + weights = Qwen4ExpTransformerLayerWeight(1, torch.bfloat16, config, Quantcfg(config)) + weights.load_hf_weights(dict(source)) + weights.verify_load() + infer = Qwen4ExpTransformerLayerInfer(1, config) + channels = 512 + conv = torch.zeros(2, channels, 3, device="cuda", dtype=torch.bfloat16) + recurrent = torch.zeros(2, 2, 128, 128, device="cuda", dtype=torch.float32) + state = Namespace( + is_prefill=False, + need_dp_prefill_balance=False, + dist_group=None, + batch_size=2, + b_req_idx=torch.tensor([0, 1], device="cuda", dtype=torch.int32), + att_state_role="main", + att_state_manager=SimpleNamespace( + get_state=lambda **kwargs: SimpleNamespace(get_layer_view=lambda **kwargs: (conv, recurrent)) + ), + req_manager=SimpleNamespace(get_mamba_cache=lambda layer: (conv, recurrent)), + ) + + class Model: + pass + + model = Model() + model.config, model.tp_world_size_ = config, 1 + model.is_mtp_draft_model = False + model.mtp_manager = SimpleNamespace(get_decode_draft_step=lambda is_draft: 0) + backend = TritonLinearAttBackend(model=model) + state.decode_att_state1 = backend.create_att_decode_state(state) + state.decode_att_state1.init_state() + cache = LexicalRequestCache(weights, 2) + ref_ngram, ref_ple = cache.ngram.clone(), cache.convolution.clone() + ref_conv, ref_recurrent = conv.clone(), recurrent.clone() + ones = torch.ones(2, device="cuda", dtype=torch.int32) + metadata = PackedBatchMetadata.prepare(state.b_req_idx, ones, ones, 2, 7, 4) + + def run(residual, tokens): + batch = cache.prepare(metadata) + state.input_ids, state.qwen4_exp_ple = tokens, {1: batch} + output = infer.token_forward(residual, state, weights) + cache.commit(metadata, batch) + return output + + if capture: + residual = torch.randn(2, 64, device="cuda", dtype=torch.bfloat16) + tokens = torch.zeros(2, device="cuda", dtype=torch.int64) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + run(residual, tokens) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_output = run(residual, tokens) + conv.zero_() + recurrent.zero_() + for request in range(2): + cache.reset(request) + for ids in [[7, 11], [63, 15], [19, 23]]: + if capture: + residual.normal_() + tokens.copy_(torch.tensor(ids, device="cuda", dtype=torch.int64)) + graph.replay() + actual = captured_output + else: + residual = torch.randn(2, 64, device="cuda", dtype=torch.bfloat16) + tokens = torch.tensor(ids, device="cuda", dtype=torch.int64) + actual = run(residual, tokens) + expected, ref_ngram, ref_ple, ref_conv, ref_recurrent = reference_step( + residual, tokens, source, config, ref_ngram, ref_ple, ref_conv, ref_recurrent + ) + torch.testing.assert_close(actual, expected, atol=0.016, rtol=0.016) + torch.testing.assert_close(cache.ngram, ref_ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution, ref_ple, atol=0.016, rtol=0.016) + torch.testing.assert_close(conv, ref_conv, atol=0.016, rtol=0.016) + + +def test_native_packed_gdn_ple_prefill_matches_recurrent_reference(config): + from lightllm.models.qwen3next.model import Qwen3NextTpPartModel + + torch.manual_seed(1037) + + Qwen3NextTpPartModel._init_triton(None) + source = checkpoint(config, 1) + weights = Qwen4ExpTransformerLayerWeight(1, torch.bfloat16, config, Quantcfg(config)) + weights.load_hf_weights(dict(source)) + weights.verify_load() + infer = Qwen4ExpTransformerLayerInfer(1, config) + conv = torch.zeros(2, 512, 3, device="cuda", dtype=torch.bfloat16) + recurrent = torch.zeros(2, 2, 128, 128, device="cuda", dtype=torch.float32) + starts = torch.tensor([0, 3, 5], device="cuda", dtype=torch.int32) + tokens = torch.tensor([7, 63, 19, 11, 23], device="cuda", dtype=torch.int64) + residual = torch.randn(5, 64, device="cuda", dtype=torch.bfloat16) + batch = PLEBatch(starts, weights.ngram_hasher.component.empty_history(2), weights.ple.component.empty_history(2)) + state = Namespace( + is_prefill=True, + need_dp_prefill_balance=False, + dist_group=None, + batch_size=2, + b_req_idx=torch.tensor([0, 1], device="cuda", dtype=torch.int32), + b1_cu_q_seq_len=starts, + b_ready_cache_len=torch.zeros(2, device="cuda", dtype=torch.int32), + att_state_role="main", + att_state_manager=SimpleNamespace( + get_state=lambda **kwargs: SimpleNamespace(get_layer_view=lambda **kwargs: (conv, recurrent)) + ), + req_manager=SimpleNamespace(get_mamba_cache=lambda layer: (conv, recurrent)), + input_ids=tokens, + qwen4_exp_ple={1: batch}, + ) + + class Model: + pass + + model = Model() + model.config, model.tp_world_size_ = config, 1 + model.is_mtp_draft_model = False + model.mtp_manager = SimpleNamespace(get_decode_draft_step=lambda is_draft: 0) + backend = TritonLinearAttBackend(model=model) + state.prefill_att_state1 = backend.create_att_prefill_state(state) + state.prefill_att_state1.init_state() + actual = infer.context_forward(residual, state, weights) + next_ngram, next_ple = batch.histories_after(weights) + expected = [] + for request, (start, end) in enumerate([(0, 3), (3, 5)]): + ng = weights.ngram_hasher.component.empty_history(1) + ple = weights.ple.component.empty_history(1) + cv = torch.zeros(1, 512, 3, device="cuda", dtype=torch.bfloat16) + rc = torch.zeros(1, 2, 128, 128, device="cuda", dtype=torch.float32) + for token in range(start, end): + out, ng, ple, cv, rc = reference_step( + residual[token : token + 1], tokens[token : token + 1], source, config, ng, ple, cv, rc + ) + expected.append(out) + torch.testing.assert_close(next_ngram[request : request + 1], ng, atol=0, rtol=0) + torch.testing.assert_close(next_ple[request : request + 1], ple, atol=0.016, rtol=0.016) + torch.testing.assert_close(actual, torch.cat(expected), atol=0.016, rtol=0.016) + + +def test_native_ple_empty_packed_batch_preserves_histories(config): + source = checkpoint(config, 1) + weights = Qwen4ExpTransformerLayerWeight(1, torch.bfloat16, config, Quantcfg(config)) + weights.load_hf_weights(dict(source)) + weights.verify_load() + starts = torch.zeros(3, device="cuda", dtype=torch.int32) + ngram = weights.ngram_hasher.component.empty_history(2) + conv = weights.ple.component.empty_history(2) + batch = PLEBatch(starts, ngram, conv) + update = batch.apply( + torch.empty(0, 64, device="cuda", dtype=torch.bfloat16), + torch.empty(0, device="cuda", dtype=torch.int64), + weights, + 1, + None, + ) + assert update.shape == (0, 64) + next_ngram, next_conv = batch.histories_after(weights) + torch.testing.assert_close(next_ngram, ngram, atol=0, rtol=0) + torch.testing.assert_close(next_conv, conv, atol=0, rtol=0) diff --git a/unit_tests/models/qwen4_exp/test_lm_head.py b/unit_tests/models/qwen4_exp/test_lm_head.py new file mode 100644 index 0000000000..27300687aa --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_lm_head.py @@ -0,0 +1,45 @@ +import pytest +import torch + +from lightllm.models.qwen4_exp.layer_weights.lm_head_weight import Qwen4ExpLMHeadWeight + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("rows", [0, 1, 3, 8, 9]) +def test_head_shape_precision_and_explicit_output(dtype, rows): + torch.manual_seed(3107) + head = Qwen4ExpLMHeadWeight(128, 257, "head.weight", dtype) + head.weight.normal_(std=0.02) + hidden = torch.randn(rows, 128, device="cuda", dtype=dtype) + result = head(hidden.T) + assert result.shape == (257, rows) and result.dtype == dtype and result.is_contiguous() + expected = (head.weight.float() @ hidden.float().T).to(dtype) + tolerance = max(1e-6, torch.finfo(dtype).eps) + torch.testing.assert_close(result, expected, atol=tolerance, rtol=tolerance) + out = torch.empty(257, rows * 2, device="cuda", dtype=dtype)[:, ::2] + assert head(hidden.T, out=out) is out + torch.testing.assert_close(out, result, atol=0, rtol=0) + + +def test_trained_geometry_batch_invariance_and_graph_replay(): + torch.manual_seed(3107) + head = Qwen4ExpLMHeadWeight(2560, 62080, "head.weight", torch.bfloat16) + head.weight.normal_(std=0.02) + hidden = torch.randn(24, 2560, device="cuda", dtype=torch.bfloat16) + singles = torch.cat([head(row[:, None]) for row in hidden], dim=1) + for chunk in (2, 3, 5, 7, 8, 9, 15, 24): + result = torch.cat([head(hidden[start : start + chunk].T) for start in range(0, len(hidden), chunk)], dim=1) + torch.testing.assert_close(result, singles, atol=0, rtol=0) + + inputs = hidden[:3].clone() + for _ in range(3): + head(inputs.T) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = head(inputs.T) + inputs.add_(0.125) + graph.replay() + torch.testing.assert_close(captured, head(inputs.T), atol=0, rtol=0) + graph.reset() diff --git a/unit_tests/models/qwen4_exp/test_mtp.py b/unit_tests/models/qwen4_exp/test_mtp.py new file mode 100644 index 0000000000..48c627a3dc --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_mtp.py @@ -0,0 +1,148 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.hidden_collector import FinalHiddenCollector +from lightllm.models.draft_registry import get_draft_model_class +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel +from lightllm.models.qwen4_exp.mtp.model import Qwen4ExpMTPModel +from lightllm.models.qwen4_exp.mtp.infer import Qwen4ExpMTPPreLayerInfer +from .helpers import checkpoint, prepost_checkpoint, add_hyper, rand +from .runtime import state_for +from .mtp_reference import reference_fusion, reference_draft + + +def test_public_draft_registry(): + for model_type in ("qwen4_exp", "qwen4_exp_text"): + assert get_draft_model_class({"model_type": model_type}, "eagle_with_att") is Qwen4ExpMTPModel + + +def build_draft(runtime, config): + main = runtime.model + source = prepost_checkpoint(config) + main.pre_post_weight = main.pre_and_post_weight_class(torch.bfloat16, config) + main.pre_post_weight.load_hf_weights(dict(source)) + main.pre_post_weight.verify_load() + main.layers_infer = [None] * config["n_layer"] + draft = Qwen4ExpMTPModel.__new__(Qwen4ExpMTPModel) + draft.__dict__.update( + main_model=main, + mtp_previous_draft_models=[], + args=main.args, + weight_dir_=main.args.model_dir, + finetune_config=None, + tp_world_size_=1, + tp_rank_=0, + torch_memory_saver=main.torch_memory_saver, + data_type=torch.bfloat16, + quant_cfg=runtime.weight.quant_cfg, + max_req_num=main.max_req_num, + prefill_graph=None, + ) + draft._init_config() + draft._init_weights() + from .resource_helpers import bind_states + + draft.bind_req_manager(main.req_manager) + draft._init_infer_layer() + main.att_state_manager._state_registration_finished = False + draft.prepare_attention_states() + runtime.mem = bind_states(main) + main.att_state_manager.get_state.cache_clear() + draft._init_custom() + draft._init_infer_layer() + draft._init_some_value() + draft._init_att_backend() + draft._init_att_backend1() + draft.mtp_manager = SimpleNamespace(get_decode_draft_step=lambda is_draft: 0) + for name, value in checkpoint(config, 3).items(): + source[name.replace("model.language_model.layers.3.", "mtp.layers.0.", 1)] = value + add_hyper(source, "mtp.hyper_connection_mixer", config, False) + source.update( + { + "mtp.pre_fc_norm_embedding.weight": rand(config["hidden_size"]), + "mtp.pre_fc_norm_hidden.weight": rand(config["hidden_size"] * config["hc_count"]), + "mtp.fc_embedding.weight": rand(config["hidden_size"], config["hidden_size"]), + "mtp.fc_hidden.weight": rand(config["hidden_size"], config["hidden_size"]), + } + ) + + for name in reversed(list(source)): + draft.pre_post_weight.load_hf_weights({name: source[name]}) + draft.trans_layers_weight[0].load_hf_weights({name: source[name]}) + draft.pre_post_weight.verify_load() + draft.trans_layers_weight[0].verify_load() + return draft, source + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("runtime", [{"mtp_step": 2, "mtp_mode": "eagle_with_att"}], indirect=True) +def test_native_complete_draft_prefill_cached_decode_and_full_hidden(runtime, config): + torch.manual_seed(2307) + Qwen3NextTpPartModel._init_triton(None) + draft, source = build_draft(runtime, config) + assert draft.req_manager is runtime.req and draft.att_state_manager is runtime.model.att_state_manager + assert draft.pre_post_weight.wte_weight_ is runtime.model.pre_post_weight.wte_weight_ + assert draft.pre_post_weight.lm_head_weight_ is runtime.model.pre_post_weight.lm_head_weight_ + assert draft.layers_infer[0].layer_num_ == 4 + assert not hasattr(draft.trans_layers_weight[0], "ple") + for request in (1, 3): + runtime.req.init_linear_att_state(SimpleNamespace(req_idx=request)) + owner = SimpleNamespace(model=draft, req=runtime.req, mem=runtime.mem) + tokens = torch.tensor([7, 11, 13, 17, 19], device="cuda", dtype=torch.int64) + hidden = torch.randn(5, 64, device="cuda", dtype=torch.bfloat16) + hidden_before = hidden.clone() + state = state_for(owner, [1, 3], [3, 2], tokens.tolist(), ready=[0, 0]) + state.mtp_draft_input_hiddens = hidden + state.hidden_collector = FinalHiddenCollector() + assert not state.qwen4_exp_ple + actual = draft._context_forward(state) + positions = torch.tensor([0, 1, 2, 0, 1], device="cuda") + cos, sin = draft._cos_cached[positions].repeat(1, 2), draft._sin_cached[positions].repeat(1, 2) + multi, logits = reference_draft(tokens, hidden, source, config, cos, sin, [3, 2]) + torch.testing.assert_close(actual.logits, logits[[2, 4]], atol=0.025, rtol=0.025) + torch.testing.assert_close(actual.mtp_collector.spec_hidden, multi, atol=0.016, rtol=0.016) + torch.testing.assert_close(hidden, hidden_before, atol=0, rtol=0) + fusion = draft.pre_post_weight.fusion.component + embeddings = source["model.language_model.embed_tokens.weight"].cuda()[tokens] + torch.testing.assert_close( + fusion(embeddings, hidden), reference_fusion(embeddings, hidden, source, config), atol=0.016, rtol=0.016 + ) + with pytest.raises(ValueError, match="full pre-mixer"): + fusion(embeddings, hidden[:, :32].contiguous()) + + image_pre = Qwen4ExpMTPPreLayerInfer(dict(config, image_token_id=60)) + virtual = torch.tensor([7, 1000000, 1000001], device="cuda", dtype=torch.int64) + image_state = SimpleNamespace( + multimodal_params=[{"images": [{"token_id": 1000000, "token_num": 2}]}], + mtp_draft_input_hiddens=hidden[:3], + dist_group=None, + ) + canonical = torch.tensor([7, 60, 60], device="cuda") + image_input = image_pre.context_forward(virtual, image_state, draft.pre_post_weight) + image_embedding = source["model.language_model.embed_tokens.weight"].cuda()[canonical] + torch.testing.assert_close( + image_input, reference_fusion(image_embedding, hidden[:3], source, config), atol=0.016, rtol=0.016 + ) + assert virtual.tolist() == [7, 1000000, 1000001] + + state = state_for(owner, [3, 1], [3, 4], [23, 29]) + new_hidden = torch.randn(2, 64, device="cuda", dtype=torch.bfloat16) + state.mtp_draft_input_hiddens = new_hidden + state.hidden_collector = FinalHiddenCollector() + actual = draft._token_forward(state) + full_tokens = torch.tensor([7, 11, 13, 29, 17, 19, 23], device="cuda") + full_hidden = torch.cat([hidden[:3], new_hidden[1:2], hidden[3:], new_hidden[:1]]) + positions = torch.tensor([0, 1, 2, 3, 0, 1, 2], device="cuda") + multi, logits = reference_draft( + full_tokens, + full_hidden, + source, + config, + draft._cos_cached[positions].repeat(1, 2), + draft._sin_cached[positions].repeat(1, 2), + [4, 3], + ) + torch.testing.assert_close(actual.logits, logits[[6, 3]], atol=0.025, rtol=0.025) + torch.testing.assert_close(actual.mtp_collector.spec_hidden, multi[[6, 3]], atol=0.016, rtol=0.016) diff --git a/unit_tests/models/qwen4_exp/test_multimodal.py b/unit_tests/models/qwen4_exp/test_multimodal.py new file mode 100644 index 0000000000..6cb7bc2b10 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_multimodal.py @@ -0,0 +1,213 @@ +from copy import deepcopy +import json +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from lightllm.models.qwen4_exp.configuration import normalize_model_config +from lightllm.models.qwen4_exp.multimodal import canonical_image_ids +from lightllm.server.router.model_infer.infer_batch import g_infer_context +from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector +from .runtime import assemble_model, state_for +from .reference import reference_model + + +def test_model_types_resolve_through_public_registry(): + from lightllm.models.registry import ModelRegistry + from lightllm.models.qwen4_exp.model import Qwen4ExpTpPartModel + + for model_type in ("qwen4_exp", "qwen4_exp_text"): + assert ModelRegistry.get_model_class({"model_type": model_type}) is Qwen4ExpTpPartModel + + +@pytest.mark.parametrize("model_type", ["qwen4_exp", "qwen4_exp_text"]) +def test_startup_detects_hybrid_cache_and_optional_vision(model_type, tmp_path): + from lightllm.utils.config_utils import is_linear_att_mixed_model, has_vision_module, auto_set_fused_shared_experts + + config = {"model_type": model_type, "architectures": ["Qwen4ExpForConditionalGeneration"]} + if model_type == "qwen4_exp": + config["vision_config"] = {"hidden_size": 32} + (tmp_path / "config.json").write_text(json.dumps(config)) + assert is_linear_att_mixed_model(str(tmp_path)) + assert has_vision_module(str(tmp_path)) == (model_type == "qwen4_exp") + args = SimpleNamespace(model_dir=str(tmp_path), enable_ep_moe=False, enable_fused_shared_experts=False) + auto_set_fused_shared_experts(args) + assert args.enable_fused_shared_experts + + +def test_tokenizer_uses_native_processor_and_virtual_image_ranges(monkeypatch): + from lightllm.server import tokenizer as tokenizer_module + from lightllm.models.qwen4_exp.image_processing import Qwen4ExpImageProcessor + from lightllm.models.qwen4_exp.tokenizer import Qwen4ExpTokenizer + + config = dict( + model_type="qwen4_exp", + architectures=["Qwen4ExpForConditionalGeneration"], + vision_config={}, + vision_start_token_id=61, + vision_end_token_id=62, + image_token_id=60, + ) + monkeypatch.setattr(tokenizer_module.AutoTokenizer, "from_pretrained", lambda *args, **kwargs: object()) + monkeypatch.setattr(tokenizer_module.PretrainedConfig, "get_config_dict", lambda *args: (config, {})) + processor = SimpleNamespace(min_pixels=16, max_pixels=64, patch_size=2, merge_size=2) + monkeypatch.setattr(Qwen4ExpImageProcessor, "from_pretrained", lambda *args: processor) + tokenizer = tokenizer_module.get_tokenizer("checkpoint") + assert type(tokenizer) is Qwen4ExpTokenizer + image = SimpleNamespace(token_id=1000000, token_num=2, start_idx=None) + actual = tokenizer.encode([7, 61, 60, 62, 11], SimpleNamespace(images=[image])) + assert actual == [7, 61, 1000000, 1000001, 62, 11] + assert image.start_idx == 2 + assert tokenizer.image_processor is processor + + +@pytest.mark.parametrize("image_count", [0, 1, 2]) +@pytest.mark.parametrize("qwen4", [False, True]) +def test_token_count_hook_matches_qwen4_expansion_and_preserves_legacy(image_count, qwen4): + from lightllm.models.qwen2_vl.model import QWen2VLTokenizer + from lightllm.models.qwen4_exp.tokenizer import Qwen4ExpTokenizer + from lightllm.server.httpserver.manager import HttpServerManager + + cls = Qwen4ExpTokenizer if qwen4 else QWen2VLTokenizer + processor = SimpleNamespace(min_pixels=16, max_pixels=64, patch_size=2, merge_size=2) + tokenizer = cls( + tokenizer=object(), + image_processor=processor, + model_cfg=dict(vision_start_token_id=61, vision_end_token_id=62, image_token_id=60), + ) + images = [] + for index in range(image_count): + image = SimpleNamespace(image_w=4 * (index + 1), image_h=4, token_id=1000 + index * 100) + image.token_num = tokenizer.get_image_token_length(image) + images.append(image) + multimodal = SimpleNamespace(images=images, audios=[]) + prompt = [7] + [61, 60, 62] * image_count + [11] + expanded = tokenizer.encode(prompt, multimodal) + manager = HttpServerManager.__new__(HttpServerManager) + manager.tokenizer = tokenizer + checked = [] + manager._assert_image_token_count = checked.append + actual = manager.tokens(prompt, multimodal, SimpleNamespace()) + assert checked == [image.token_num for image in images] + assert actual == len(expanded) + + +def test_plain_text_counting_fallback_is_unchanged(): + from lightllm.server.httpserver.manager import HttpServerManager + + manager = HttpServerManager.__new__(HttpServerManager) + manager.tokenizer = SimpleNamespace(encode=lambda prompt, *_args, **_kwargs: prompt) + assert manager.tokens([7, 11], SimpleNamespace(images=[], audios=[]), SimpleNamespace()) == 2 + + +def test_token_count_hook_cannot_bypass_image_limits(): + from lightllm.server.httpserver.manager import HttpServerManager + from lightllm.models.qwen4_exp.tokenizer import Qwen4ExpTokenizer + + tokenizer = Qwen4ExpTokenizer( + tokenizer=object(), + image_processor=SimpleNamespace(min_pixels=16, max_pixels=64, patch_size=2, merge_size=2), + model_cfg=dict(vision_start_token_id=61, vision_end_token_id=62, image_token_id=60), + ) + manager = HttpServerManager.__new__(HttpServerManager) + manager.tokenizer = tokenizer + + def reject(_count): + raise ValueError("image token limit") + + manager._assert_image_token_count = reject + with pytest.raises(ValueError, match="image token limit"): + manager.tokens( + [61, 60, 62], + SimpleNamespace(images=[SimpleNamespace(image_w=4, image_h=4)], audios=[]), + SimpleNamespace(), + ) + + +def test_image_processor_loads_real_config_without_runtime_enum_serialization(tmp_path): + from lightllm.models.qwen4_exp.image_processing import Qwen4ExpImageProcessor + + (tmp_path / "preprocessor_config.json").write_text( + json.dumps( + dict( + size={"shortest_edge": 16, "longest_edge": 64}, + patch_size=2, + merge_size=2, + image_mean=[0.5] * 3, + image_std=[0.5] * 3, + image_processor_type="Qwen2VLImageProcessorFast", + ) + ) + ) + processor = Qwen4ExpImageProcessor.from_pretrained(str(tmp_path)) + assert processor.min_pixels == 16 and processor.max_pixels == 64 + assert processor.patch_size == processor.merge_size == 2 + assert processor.image_mean == processor.image_std == [0.5] * 3 + + +def test_normalize_outer_and_text_config_preserves_checkpoint(config): + text = dict(config, rope_parameters={"rope_theta": 1000, "partial_rotary_factor": 0.5}) + text.pop("n_layer") + outer = {"text_config": text, "vision_config": {"hidden_size": 32}, "image_token_id": 60} + original = deepcopy(outer) + actual, vision = normalize_model_config(outer) + standalone, no_vision = normalize_model_config(text) + assert actual["n_layer"] == standalone["n_layer"] == 4 + assert actual["image_token_id"] == 60 + assert actual["rope_theta"] == standalone["rope_theta"] == 1000 + assert vision == outer["vision_config"] and no_vision is None + assert outer == original + customized, _ = normalize_model_config(text, SimpleNamespace(vocab_size=128)) + assert customized["vocab_size"] == 128 and text["vocab_size"] == 64 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) +def test_image_ranges_preserve_text_and_physical_ids(dtype): + second = 2 ** 34 if dtype == torch.int64 else 2000000 + tokens = torch.tensor( + [7, 60, 999999, 1000000, 1000001, 1000002, second + 4, second + 5, 63], device="cuda", dtype=dtype + ) + starts = torch.tensor([1000000, second], device="cuda", dtype=torch.int64) + lengths = torch.tensor([2, 5], device="cuda", dtype=torch.int64) + before = tokens.clone() + actual = canonical_image_ids(tokens, starts, lengths, 60) + expected = torch.tensor([7, 60, 999999, 60, 60, 1000002, 60, second + 5, 63], device="cuda", dtype=dtype) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + torch.testing.assert_close(tokens, before, atol=0, rtol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = canonical_image_ids(tokens, starts, lengths, 60) + tokens[3:5] = 11 + graph.replay() + expected[3:5] = 11 + torch.testing.assert_close(captured, expected, atol=0, rtol=0) + graph.reset() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_image_cache_ids_are_canonicalized_before_ple_commit(runtime, config, monkeypatch): + config["image_token_id"] = 60 + model, source = assemble_model(runtime, config) + runtime.req.init_linear_att_state(SimpleNamespace(req_idx=1)) + image_cache = torch.randn(4, 1, 32, dtype=torch.bfloat16).pin_memory() + monkeypatch.setattr(g_infer_context, "cpu_embed_cache_client", SimpleNamespace(cpu_embed_cache_tensor=image_cache)) + virtual = [7, 1000000, 1000001, 11] + canonical = torch.tensor([7, 60, 60, 11], device="cuda", dtype=torch.int64) + state = state_for(runtime, [1], [4], virtual, ready=[0]) + state.hidden_collector = NoopHiddenCollector() + state.multimodal_params = [ + {"images": [{"token_id": 1000000, "token_num": 2, "start_index_in_embed_cache": 1}], "audios": []} + ] + actual = model._context_forward(state) + torch.testing.assert_close(state.input_ids, torch.tensor(virtual, device="cuda"), atol=0, rtol=0) + torch.testing.assert_close(state.qwen4_exp_ple[1].input_tokens, canonical, atol=0, rtol=0) + torch.testing.assert_close(runtime.req.lexical_caches[1].ngram[1], canonical[-2:], atol=0, rtol=0) + embeddings = F.embedding(canonical.cpu(), source["model.language_model.embed_tokens.weight"]) + embeddings[1:3] = image_cache[1:3, 0] + positions = torch.arange(4, device="cuda") + cos, sin = model._cos_cached[positions].repeat(1, 2), model._sin_cached[positions].repeat(1, 2) + expected = reference_model(canonical, cos, sin, [4], source, config, input_embeddings=embeddings.cuda())[-1:] + torch.testing.assert_close(actual.logits, expected, atol=0.025, rtol=0.025) diff --git a/unit_tests/models/qwen4_exp/test_norm.py b/unit_tests/models/qwen4_exp/test_norm.py new file mode 100644 index 0000000000..50ccc871c0 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_norm.py @@ -0,0 +1,48 @@ +import pytest +import torch + +from lightllm.models.qwen4_exp.layer_weights.norm_weight import Qwen4ExpGatedNormWeight + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("activation", ["sigmoid", "silu"]) +def test_native_gdn_norm_rounding(dtype, activation): + torch.manual_seed(837) + native = Qwen4ExpGatedNormWeight(128, "linear.norm.weight", dtype, activation) + weight = 1 + torch.randn(128, dtype=dtype) * 0.1 + native.load_hf_weights({"linear.norm.weight": weight}) + values = torch.randn(19, 128, device="cuda", dtype=dtype) + gate = torch.randn_like(values) * 3 + normalized = values.float() * torch.rsqrt(values.float().square().mean(-1, keepdim=True) + 1e-6) + weighted = native.weight * normalized.to(dtype) + gate_activation = gate.float().sigmoid() + if activation == "silu": + gate_activation *= gate.float() + expected = (weighted * gate_activation).to(dtype) + out = torch.empty_like(values) + actual = native(values, gate, 1e-6, out=out) + assert actual is out + tolerance = 2e-6 if dtype == torch.float32 else torch.finfo(dtype).eps * 2 + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + + +@pytest.mark.parametrize("activation", ["sigmoid", "silu"]) +def test_gdn_norm_accepts_strided_gate(activation): + from lightllm.models.qwen4_exp.triton_kernel.gated_norm import gated_norm + + torch.manual_seed(123) + tokens, heads, dim = 5, 4, 128 + values = torch.randn(tokens * heads, dim, device="cuda", dtype=torch.bfloat16) + packed = torch.randn(tokens, heads * dim + 256, device="cuda", dtype=torch.bfloat16) + gate = packed[:, 128 : 128 + heads * dim].view(tokens, heads, dim) + weight = torch.randn(dim, device="cuda", dtype=torch.bfloat16) + + actual = gated_norm(values, gate, weight, 1e-6, activation) + normalized = values.float() * torch.rsqrt(values.float().square().mean(-1, keepdim=True) + 1e-6) + expected = normalized * weight.float() + gate = gate.reshape(-1, dim).float() + gate = gate.sigmoid() if activation == "sigmoid" else torch.nn.functional.silu(gate) + expected = (expected * gate).to(values.dtype) + torch.testing.assert_close(actual, expected, atol=0.008, rtol=0.008) diff --git a/unit_tests/models/qwen4_exp/test_pinned_memory.py b/unit_tests/models/qwen4_exp/test_pinned_memory.py new file mode 100644 index 0000000000..cffc72cf36 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_pinned_memory.py @@ -0,0 +1,112 @@ +import gc +import weakref + +import pytest +import torch +import triton +import triton.language as tl + +from lightllm.models.qwen4_exp import pinned_memory + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@triton.jit +def read_ends(Source, Output, SIZE: tl.constexpr): + row = tl.arange(0, 2) + position = tl.where(row == 0, 0, SIZE - 1) + tl.store(Output + row, tl.load(Source + position)) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_registered_snapshot_views_graph_and_lifetime(dtype, monkeypatch): + monkeypatch.setattr(pinned_memory, "_LARGE_BYTES", 1) + monkeypatch.setattr(pinned_memory, "_CHUNK_BYTES", 4096) + value = pinned_memory.pinned_zeros((6144,), dtype=dtype) + assert value.is_pinned() and value.is_contiguous() + assert torch.count_nonzero(value) == 0 + value[0], value[-1] = 3, 7 + output = torch.empty(2, device="cuda", dtype=torch.float32) + read_ends[(1,)](value, output, value.numel()) + torch.testing.assert_close(output, torch.tensor([3, 7], device="cuda", dtype=torch.float32), atol=0, rtol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + read_ends[(1,)](value, output, value.numel()) + value[-1] = -7 + graph.replay() + torch.testing.assert_close(output, torch.tensor([3, -7], device="cuda", dtype=torch.float32), atol=0, rtol=0) + graph.reset() + owner = weakref.ref(value._qwen4_pinned_allocation) + view = value[4096:] + del value + gc.collect() + assert owner() is not None and view.is_pinned() + read_ends[(1,)](view, output, view.numel()) + torch.testing.assert_close(output, torch.tensor([0, -7], device="cuda", dtype=torch.float32), atol=0, rtol=0) + del view + gc.collect() + assert owner() is None + + +def test_partial_registration_failure_releases_registered_ranges(monkeypatch): + monkeypatch.setattr(pinned_memory, "_LARGE_BYTES", 1) + monkeypatch.setattr(pinned_memory, "_CHUNK_BYTES", 4096) + runtime = torch.cuda.cudart() + registered, released = [], [] + + class FailingRuntime: + def cudaHostRegister(self, address, size, flags): + if registered: + return 1 + status = runtime.cudaHostRegister(address, size, flags) + assert int(status) == 0 + registered.append(address) + return status + + def cudaHostUnregister(self, address): + released.append(address) + return runtime.cudaHostUnregister(address) + + def cudaGetErrorString(self, status): + return "intentional registration failure" + + monkeypatch.setattr(torch.cuda, "cudart", lambda: FailingRuntime()) + + def allocate(): + try: + pinned_memory.pinned_zeros((6144,), dtype=torch.float32) + except RuntimeError as error: + assert "intentional registration failure" in str(error) + else: + pytest.fail("Expected registration failure") + + allocate() + gc.collect() + assert registered and released == registered + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_snapshot_dma_restore_never_crosses_registration(dtype, monkeypatch): + monkeypatch.setattr(pinned_memory, "_LARGE_BYTES", 1) + monkeypatch.setattr(pinned_memory, "_CHUNK_BYTES", 16384) + value = pinned_memory.pinned_zeros((7, 3, 512), dtype=dtype) + cell_bytes = value[0].numel() * value.element_size() + starts = value._qwen4_pinned_allocation.addresses + assert len(starts) > 1 + assert all((address - value.data_ptr()) % cell_bytes == 0 for address in starts) + ends = starts[1:] + [value.data_ptr() + value.numel() * value.element_size()] + assert all(end - start <= 16384 for start, end in zip(starts, ends)) + destination = torch.empty((3, 5, 512), dtype=dtype, device="cuda") + for index in range(value.shape[0]): + value[index].fill_(index + 1) + + destination[:, 1, ...] = value[index] + torch.cuda.synchronize() + torch.testing.assert_close(destination[:, 1].cpu(), value[index], atol=0, rtol=0) + + +def test_snapshot_larger_than_registration_bound_is_rejected(monkeypatch): + monkeypatch.setattr(pinned_memory, "_LARGE_BYTES", 1) + monkeypatch.setattr(pinned_memory, "_CHUNK_BYTES", 16384) + with pytest.raises(ValueError, match="one bounded host registration"): + pinned_memory.pinned_zeros((2, 8192), dtype=torch.float32) diff --git a/unit_tests/models/qwen4_exp/test_reduction.py b/unit_tests/models/qwen4_exp/test_reduction.py new file mode 100644 index 0000000000..2eaf564f06 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_reduction.py @@ -0,0 +1,104 @@ +import datetime +import json +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + + +def _check_reductions(rank, rendezvous, cuda=False): + import os + + from lightllm.models.qwen4_exp.layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils.envs_utils import get_env_start_args + + args = StartArgs() + args.enable_tpsp_mix_mode = False + os.environ["LIGHTLLM_START_ARGS"] = json.dumps(vars(args)) + get_env_start_args.cache_clear() + device = torch.device("cuda", rank) if cuda else torch.device("cpu") + if cuda: + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl" if cuda else "gloo", + init_method="file://" + rendezvous, + rank=rank, + world_size=4, + timeout=datetime.timedelta(seconds=60), + ) + try: + layer = Qwen4ExpTransformerLayerInfer.__new__(Qwen4ExpTransformerLayerInfer) + layer.tp_world_size_ = 4 + layer.tp_rank_ = rank + state = SimpleNamespace(dist_group=dist.group.WORLD) + + contribution = (256.0, 1.0, -256.0, 0.25)[rank] + for dtype in (torch.bfloat16, torch.float16, torch.float32): + for rows in (1, 3, 8): + value = torch.full((rows, 16), contribution, dtype=dtype, device=device) + value[:, 1::2].neg_() + expected = torch.full_like(value, 1.25) + expected[:, 1::2].neg_() + for reduce in (layer._reduce_linear_output, layer._tpsp_reduce): + actual = reduce(value.clone(), state) + assert actual.dtype == dtype + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + midpoint_parts = (1.0, 2 ** -8, 2 ** -24, 2 ** -24) + for rows in (1, 3, 256, 512): + value = torch.full((rows, 16), midpoint_parts[rank], dtype=torch.bfloat16, device=device) + expected = torch.full_like(value, 1.0078125) + for reduce in (layer._reduce_linear_output, layer._tpsp_reduce): + actual = reduce(value.clone(), state) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + value = torch.full((16, 3), contribution, dtype=torch.bfloat16, device=device).t() + assert layer._reduce_linear_output(value, state) is value + torch.testing.assert_close(value, torch.full_like(value, 1.25), atol=0, rtol=0) + + args.enable_tpsp_mix_mode = True + os.environ["LIGHTLLM_START_ARGS"] = json.dumps(vars(args)) + get_env_start_args.cache_clear() + for rows in (4, 12): + common = torch.arange(rows * 16, device=device, dtype=torch.float32).reshape(rows, 16) / 16 + value = (common + rank / 4).to(torch.bfloat16) + expected = (common * 4 + 1.5).to(value.dtype).chunk(4)[rank] + actual = layer._tpsp_reduce(value, state) + assert actual.shape == (rows // 4, 16) and actual.dtype == value.dtype + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + if cuda: + for _ in range(3): + layer._tpsp_reduce(value, state) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + result = layer._tpsp_reduce(value, state) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(result, expected, atol=0, rtol=0) + graph.reset() + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason="Gloo is required") +def test_tp_and_sequence_parallel_accumulation(tmp_path): + mp.spawn(_check_reductions, args=(str(tmp_path / "rendezvous"),), nprocs=4, join=True) + + +@pytest.mark.skipif(torch.cuda.device_count() < 4 or not dist.is_nccl_available(), reason="Four CUDA GPUs are required") +def test_cuda_tp_and_sequence_parallel_accumulation(tmp_path): + mp.spawn(_check_reductions, args=(str(tmp_path / "nccl-rendezvous"), True), nprocs=4, join=True) + + +def test_single_rank_preserves_tensor(): + from lightllm.models.qwen4_exp.layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer + + layer = Qwen4ExpTransformerLayerInfer.__new__(Qwen4ExpTransformerLayerInfer) + layer.tp_world_size_ = 1 + value = torch.tensor([[1.25]], dtype=torch.bfloat16) + assert layer._reduce_linear_output(value, None) is value + assert layer._tpsp_reduce(value, None) is value diff --git a/unit_tests/models/qwen4_exp/test_runtime.py b/unit_tests/models/qwen4_exp/test_runtime.py new file mode 100644 index 0000000000..25fe2c0f79 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_runtime.py @@ -0,0 +1,228 @@ +import json +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.qwen4_exp.layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight +from lightllm.models.qwen4_exp.model import Qwen4ExpTpPartModel +from lightllm.models.qwen4_exp.layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer +from lightllm.common.quantization import Quantcfg +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel +from .helpers import checkpoint +from .reference import reference_step, reference_layer, reference_model +from .runtime import state_for, assemble_model +from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("runtime", [{}, {"ssm_state_mode": "replay"}], ids=["native", "replay"], indirect=True) +def test_inference_state_prefill_decode_and_automatic_commit(runtime, config): + Qwen3NextTpPartModel._init_triton(None) + req, weight = runtime.req, runtime.weight + for request in (1, 3): + req.init_linear_att_state(SimpleNamespace(req_idx=request)) + state = state_for(runtime, [1, 3], [3, 2], [7, 63, 19, 11, 23], ready=[0, 0]) + assert state.qwen4_exp_metadata.starts.tolist() == [0, 3, 5, 5, 5, 5] + assert state.qwen4_exp_batch.slot_table.data_ptr() == req.req_to_token_indexs.data_ptr() + torch.testing.assert_close( + runtime.mem.qsa_key_cos[state.mem_index.long()], state.qwen4_exp_batch.cos, atol=0, rtol=0 + ) + infer = Qwen4ExpTransformerLayerInfer(1, config) + residual = torch.randn(5, 64, device="cuda", dtype=torch.bfloat16) + reference = {} + expected = [] + for request, start, end in ((1, 0, 3), (3, 3, 5)): + ngram = weight.ngram_hasher.component.empty_history(1) + ple = weight.ple.component.empty_history(1) + conv = torch.zeros(1, 512, 3, device="cuda", dtype=torch.bfloat16) + ssm = torch.zeros(1, 2, 128, 128, device="cuda", dtype=torch.float32) + for row in range(start, end): + output, ngram, ple, conv, ssm = reference_step( + residual[row : row + 1], state.input_ids[row : row + 1], runtime.source, config, ngram, ple, conv, ssm + ) + expected.append(output) + reference[request] = (ngram, ple, conv, ssm) + actual = infer.context_forward(residual, state, weight) + torch.testing.assert_close(actual, torch.cat(expected), atol=0.016, rtol=0.016) + + state = state_for(runtime, [1, 3], [4, 3], [29, 31]) + residual = torch.randn(2, 64, device="cuda", dtype=torch.bfloat16) + actual = infer.token_forward(residual, state, weight) + for row, request in enumerate((1, 3)): + expected, ngram, ple, _, _ = reference_step( + residual[row : row + 1], state.input_ids[row : row + 1], runtime.source, config, *reference[request] + ) + torch.testing.assert_close(actual[row : row + 1], expected, atol=0.016, rtol=0.016) + torch.testing.assert_close(req.lexical_caches[1].ngram[request], ngram[0], atol=0, rtol=0) + torch.testing.assert_close(req.lexical_caches[1].convolution[request], ple[0], atol=0.016, rtol=0.016) + + +def test_native_prefill_graph_state_copy_changes_request_count(runtime, config): + source = checkpoint(config, 3) + weight = Qwen4ExpTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + weight.load_hf_weights(dict(source)) + weight.verify_load() + infer = Qwen4ExpTransformerLayerInfer(3, config) + old = state_for(runtime, [1], [4], [7, 11, 13, 17], ready=[0]) + residual = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16) + for _ in range(3): + infer.context_forward(residual, old, weight) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = infer.context_forward(residual, old, weight) + pointers = [ + getattr(old.qwen4_exp_metadata, name).data_ptr() for name in old.qwen4_exp_metadata.__dataclass_fields__ + ] + before = runtime.mem.get_index_kv_buffer(old.mem_index) + new = state_for(runtime, [2, 3], [2, 2], [19, 23, 29, 31], ready=[0, 0]) + old.copy_for_prefill_cuda_graph(new) + residual.normal_() + graph.replay() + expected = reference_layer(residual, source, config, new.qwen4_exp_batch.cos, new.qwen4_exp_batch.sin, [2, 2]) + torch.testing.assert_close(output, expected, atol=0.016, rtol=0.016) + assert pointers == [ + getattr(old.qwen4_exp_metadata, name).data_ptr() for name in old.qwen4_exp_metadata.__dataclass_fields__ + ] + for name, values in runtime.mem.get_index_kv_buffer(torch.arange(4, device="cuda")).items(): + torch.testing.assert_close(values, before[name], atol=0, rtol=0, equal_nan=True) + + +def test_native_config_normalization_and_periodic_layout(runtime, config, tmp_path): + outer = dict(text_config=dict(config), vision_config={"hidden_size": 16}) + outer["text_config"]["rope_parameters"] = dict( + rope_type="default", rope_theta=10000, partial_rotary_factor=0.5, mrope_section=[1, 1, 0] + ) + (tmp_path / "config.json").write_text(json.dumps(outer)) + model = Qwen4ExpTpPartModel.__new__(Qwen4ExpTpPartModel) + model.weight_dir_, model.finetune_config, model.tp_world_size_ = str(tmp_path), None, 1 + model._init_config() + assert model.config["n_layer"] == 4 and model.config["n_embed"] == 32 + assert model.config["rope_theta"] == 10000 + assert model.num_kv_heads == 2 + invalid = dict(config, layer_types=["full_attention"] * 4) + with pytest.raises(ValueError, match="periodic"): + Qwen4ExpTpPartModel.normalize_model_config(invalid) + + +def test_native_decode_graph_replay_updates_positions_slots_and_request_order(runtime, config): + source = checkpoint(config, 3) + weight = Qwen4ExpTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + weight.load_hf_weights(dict(source)) + weight.verify_load() + infer = Qwen4ExpTransformerLayerInfer(3, config) + prefix = state_for(runtime, [1, 3], [4, 4], [7, 11, 13, 17, 19, 23, 29, 31], ready=[0, 0]) + prefix_residual = torch.randn(8, 64, device="cuda", dtype=torch.bfloat16) + infer.context_forward(prefix_residual, prefix, weight) + old = state_for(runtime, [1, 3], [5, 5], [37, 41]) + residual = torch.randn(2, 64, device="cuda", dtype=torch.bfloat16) + for _ in range(3): + infer.token_forward(residual, old, weight) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = infer.token_forward(residual, old, weight) + previous_residual = residual.clone() + previous_cos, previous_sin = old.qwen4_exp_batch.cos.clone(), old.qwen4_exp_batch.sin.clone() + new = state_for(runtime, [3, 1], [6, 6], [43, 47]) + old.copy_for_cuda_graph(new) + residual.normal_() + graph.replay() + all_residual, all_cos, all_sin = [], [], [] + for row, (start, end) in enumerate(((0, 4), (4, 8))): + incoming = 1 - row + all_residual.extend( + [prefix_residual[start:end], previous_residual[row : row + 1], residual[incoming : incoming + 1]] + ) + all_cos.extend( + [ + prefix.qwen4_exp_batch.cos[start:end], + previous_cos[row : row + 1], + new.qwen4_exp_batch.cos[incoming : incoming + 1], + ] + ) + all_sin.extend( + [ + prefix.qwen4_exp_batch.sin[start:end], + previous_sin[row : row + 1], + new.qwen4_exp_batch.sin[incoming : incoming + 1], + ] + ) + expected = reference_layer(torch.cat(all_residual), source, config, torch.cat(all_cos), torch.cat(all_sin), [6, 6]) + torch.testing.assert_close(output, expected[[11, 5]], atol=0.016, rtol=0.016) + + +@pytest.mark.parametrize("runtime", [{}, {"ssm_state_mode": "replay"}], ids=["native", "replay"], indirect=True) +def test_complete_small_model_prefill_and_cached_logits(runtime, config): + Qwen3NextTpPartModel._init_triton(None) + model, source = assemble_model(runtime, config) + for request in (1, 3): + runtime.req.init_linear_att_state(SimpleNamespace(req_idx=request)) + state = state_for(runtime, [1, 3], [3, 2], [7, 63, 19, 11, 23], ready=[0, 0]) + state.hidden_collector = NoopHiddenCollector() + output = model._context_forward(state) + expected = reference_model( + state.input_ids, state.qwen4_exp_batch.cos, state.qwen4_exp_batch.sin, [3, 2], source, config + ) + torch.testing.assert_close(output.logits, expected[[2, 4]], atol=0.025, rtol=0.025) + decode = state_for(runtime, [1, 3], [4, 3], [29, 31]) + decode.hidden_collector = NoopHiddenCollector() + output = model._token_forward(decode) + all_tokens = torch.tensor([7, 63, 19, 29, 11, 23, 31], device="cuda", dtype=torch.int64) + all_cos = torch.cat( + [ + state.qwen4_exp_batch.cos[:3], + decode.qwen4_exp_batch.cos[:1], + state.qwen4_exp_batch.cos[3:], + decode.qwen4_exp_batch.cos[1:], + ] + ) + all_sin = torch.cat( + [ + state.qwen4_exp_batch.sin[:3], + decode.qwen4_exp_batch.sin[:1], + state.qwen4_exp_batch.sin[3:], + decode.qwen4_exp_batch.sin[1:], + ] + ) + expected = reference_model(all_tokens, all_cos, all_sin, [4, 3], source, config) + torch.testing.assert_close(output.logits, expected[[3, 6]], atol=0.025, rtol=0.025) + + third = state_for(runtime, [1, 3], [5, 4], [37, 41]) + third.hidden_collector = NoopHiddenCollector() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = model._token_forward(third) + graph.replay() + third_cos, third_sin = third.qwen4_exp_batch.cos.clone(), third.qwen4_exp_batch.sin.clone() + all_tokens = torch.tensor([7, 63, 19, 29, 37, 11, 23, 31, 41], device="cuda", dtype=torch.int64) + all_cos = torch.cat([all_cos[:4], third_cos[:1], all_cos[4:], third_cos[1:]]) + all_sin = torch.cat([all_sin[:4], third_sin[:1], all_sin[4:], third_sin[1:]]) + expected = reference_model(all_tokens, all_cos, all_sin, [5, 4], source, config) + torch.testing.assert_close(captured.logits, expected[[4, 8]], atol=0.025, rtol=0.025) + fourth = state_for(runtime, [3, 1], [5, 6], [47, 43]) + fourth.hidden_collector = NoopHiddenCollector() + third.copy_for_cuda_graph(fourth) + graph.replay() + all_tokens = torch.tensor([7, 63, 19, 29, 37, 43, 11, 23, 31, 41, 47], device="cuda", dtype=torch.int64) + all_cos = torch.cat([all_cos[:5], fourth.qwen4_exp_batch.cos[1:], all_cos[5:], fourth.qwen4_exp_batch.cos[:1]]) + all_sin = torch.cat([all_sin[:5], fourth.qwen4_exp_batch.sin[1:], all_sin[5:], fourth.qwen4_exp_batch.sin[:1]]) + expected = reference_model(all_tokens, all_cos, all_sin, [6, 5], source, config) + torch.testing.assert_close(captured.logits, expected[[10, 5]], atol=0.025, rtol=0.025) + + +def test_fresh_prefill_ignores_stale_lexical_slot_without_mutating_accepted_state(runtime): + cache = runtime.req.lexical_caches[1] + cache.ngram[1].fill_(17) + cache.convolution[1].normal_() + accepted = cache.snapshot(1) + fresh = state_for(runtime, [1], [2], [7, 11], ready=[0]) + batch = fresh.qwen4_exp_ple[1] + assert torch.all(batch.ngram_history == 63) + assert torch.count_nonzero(batch.conv_history) == 0 + torch.testing.assert_close(cache.ngram[1], accepted.ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution[1], accepted.convolution, atol=0, rtol=0) + resumed = state_for(runtime, [1], [4], [13, 19], ready=[2]) + batch = resumed.qwen4_exp_ple[1] + torch.testing.assert_close(batch.ngram_history[0], accepted.ngram, atol=0, rtol=0) + torch.testing.assert_close(batch.conv_history[0], accepted.convolution, atol=0, rtol=0) diff --git a/unit_tests/models/qwen4_exp/test_sequence_parallel.py b/unit_tests/models/qwen4_exp/test_sequence_parallel.py new file mode 100644 index 0000000000..08a4aa8e67 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_sequence_parallel.py @@ -0,0 +1,87 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.qwen3_5.layer_infer.transformer_layer_infer import Qwen35TransformerLayerInfer +from lightllm.models.qwen4_exp.layer_infer import transformer_layer_infer as decoder + + +def partitioned_layer(world, rank, enabled, full): + layer = decoder.Qwen4ExpTransformerLayerInfer.__new__(decoder.Qwen4ExpTransformerLayerInfer) + layer.tp_world_size_, layer.tp_rank_ = world, rank + layer.layer_num_, layer.is_linear_attention_layer = 1, True + layer._tpsp_allgather = lambda value, state: full if enabled else value + layer._tpsp_sp_split = lambda value, state: value.chunk(world)[rank].contiguous() if enabled else value + return layer + + +@pytest.mark.parametrize("world,rank", [(2, 0), (2, 1), (4, 0), (4, 3)]) +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("method", ["context_attention_forward", "token_attention_forward"]) +def test_gdn_uses_global_rows_and_returns_local_partition(monkeypatch, world, rank, enabled, method): + full = torch.arange(64, dtype=torch.float32).reshape(8, 8) + layer = partitioned_layer(world, rank, enabled, full) + monkeypatch.setattr(decoder, "get_env_start_args", lambda: SimpleNamespace(enable_tpsp_mix_mode=enabled)) + + def ordered(value, group, size, scatter_rank=None): + total = value * size + return total if scatter_rank is None else total.chunk(size)[scatter_rank].contiguous() + + def gdn(self, value, state, weight): + assert torch.equal(value, full) + return self._reduce_linear_output(value, state) + + monkeypatch.setattr(decoder, "ordered_tp_sum", ordered) + monkeypatch.setattr(Qwen35TransformerLayerInfer, method, gdn) + value = full.chunk(world)[rank].contiguous() if enabled else full + actual = getattr(layer, method)(value, SimpleNamespace(dist_group=None), None) + expected = (full * world).chunk(world)[rank] if enabled else full * world + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("world,rank", [(2, 0), (2, 1), (4, 0), (4, 3)]) +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("prefill", [False, True]) +@pytest.mark.parametrize("commit", [None, False, True]) +def test_ple_keeps_global_metadata_and_local_residual(world, rank, enabled, prefill, commit): + full = torch.arange(64, dtype=torch.float32).reshape(8, 8) + lexical = full / 16 + layer = partitioned_layer(world, rank, enabled, full) + layer.context_attention_forward = layer.token_attention_forward = lambda value, *args: torch.zeros_like(value) + layer._ffn = lambda value, *args: torch.zeros_like(value) + tokens, metadata, batch = torch.arange(8), object(), SimpleNamespace() + calls = [] + + def apply(value, ids, weight, size, group): + assert torch.equal(value, full) and ids is tokens and size == world + return lexical + + def record(kind, observed_metadata, observed_batch): + assert observed_metadata is metadata and observed_batch is batch + calls.append(kind) + + batch.apply = apply + cache = SimpleNamespace(commit=lambda *args: record("commit", *args)) + state = SimpleNamespace( + qwen4_exp_ple={1: batch}, + input_ids=tokens, + qwen4_exp_metadata=metadata, + qwen4_exp_auto_commit=commit, + dist_group=None, + qwen4_exp_lexical_state=SimpleNamespace( + lexical_caches={1: cache}, + stage_speculative_lexical=lambda number, *args: record("stage", *args), + ), + ) + residual = SimpleNamespace(read=lambda value: (value, None), write=lambda value, update, gates: value + update) + weight = SimpleNamespace( + ple=True, + attn_hyper_connection=SimpleNamespace(component=residual), + mlp_hyper_connection=SimpleNamespace(component=residual), + ) + value = full.chunk(world)[rank].contiguous() if enabled else full + actual = layer._forward(value, state, weight, prefill) + expected = (full + lexical).chunk(world)[rank] if enabled else full + lexical + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert calls == ([] if commit is None else ["commit" if commit else "stage"]) diff --git a/unit_tests/models/qwen4_exp/test_speculative_cache.py b/unit_tests/models/qwen4_exp/test_speculative_cache.py new file mode 100644 index 0000000000..7b95cc4c84 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_speculative_cache.py @@ -0,0 +1,159 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.qwen4_exp.lexical_state import Qwen4ExpLexicalState +from lightllm.models.qwen4_exp.speculative_cache import LexicalSpeculativeCache +from lightllm.server.router.model_infer.mtp_speculative.utils import verify_mtp_tokens +from .lexical_helpers import make_cache, metadata, reference_segment +from .reference import reference_step +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel +from lightllm.models.qwen4_exp.layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("capture", [False, True]) +def test_actual_verifier_commits_only_accepted_lexical_candidates(config, capture): + torch.manual_seed(1907) + cache, source = make_cache(config) + cache.ngram[2].copy_(torch.tensor([7, 11], device="cuda")) + cache.ngram[5].copy_(torch.tensor([13, 17], device="cuda")) + cache.convolution.normal_() + untouched = cache.snapshot(7) + speculative = LexicalSpeculativeCache(cache, 4) + req = SimpleNamespace() + req.speculative_lexical_caches = {1: speculative} + req.req_sampling_params_manager = SimpleNamespace( + req_to_next_token_ids=torch.full((8, 4), -1, device="cuda", dtype=torch.int64) + ) + req.req_to_mtp_state_index = torch.zeros(8, device="cuda", dtype=torch.int32) + req.HOLD_REQUEST_ID = 7 + lexical = Qwen4ExpLexicalState.__new__(Qwen4ExpLexicalState) + lexical.request_capacity = 8 + lexical.speculative_lexical_caches = {1: speculative} + + def commit(requests, accepted, role): + assert role == "main" + lexical.commit_bounded_att_state(requests, accepted) + req.req_to_mtp_state_index[requests] = accepted - 1 + + backend = SimpleNamespace( + model=SimpleNamespace(req_manager=req, att_state_manager=SimpleNamespace(commit_bounded_att_state=commit)), + is_linear_att_mixed_model=True, + max_draft_step=3, + ) + meta = metadata([5, 5, 5, 2, 2, 7], [11, 12, 13, 11, 12, 1], [1] * 6) + tokens = torch.empty(6, device="cuda", dtype=torch.int64) + residual = torch.empty(6, 64, device="cuda", dtype=torch.bfloat16) + + def run(): + batch = cache.prepare(meta) + result = batch.apply(residual, tokens, cache.layer_weight, 1, None) + speculative.stage(meta, batch) + return result + + tokens.fill_(19) + residual.normal_() + if capture: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + result = run() + + cases = ( + ([5, 5, 5, 2, 2, 7], [2, 1], [19, 63, 23, 29, 31, 37]), + ([2, 2, 5, 5, 5, 7], [2, 3], [41, 43, 47, 63, 5, 7]), + ([5, 5, 5, 5, 2, 7], [1, 1], [11, 13, 17, 19, 23, 29]), + ) + for requests, accepts, ids in cases: + starts = [0] + [i for i in range(1, len(requests)) if requests[i] != requests[i - 1]] + counts = [end - start for start, end in zip(starts, starts[1:] + [len(requests)])] + lengths = [] + for count in counts: + lengths.extend(range(11, 11 + count)) + meta.copy_from(metadata(requests, lengths, [1] * 6)) + tokens.copy_(torch.tensor(ids, device="cuda")) + residual.normal_() + before_ngram, before_conv = cache.ngram.clone(), cache.convolution.clone() + if capture: + graph.replay() + actual = result + else: + actual = run() + torch.testing.assert_close(cache.ngram, before_ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution, before_conv, atol=0, rtol=0) + + next_ids = torch.arange(101, 107, device="cuda", dtype=torch.int64) + proposals = req.req_sampling_params_manager.req_to_next_token_ids + proposals.fill_(-1) + expected_mask = [] + for start, count, accepted in zip(starts, counts, accepts + [1]): + request = requests[start] + proposals[request, 1:accepted] = next_ids[start : start + accepted - 1] + expected_mask.extend([1] * accepted + [0] * (count - accepted)) + accept_lengths, accepted_index = verify_mtp_tokens( + backend, + next_ids, + torch.tensor(requests, device="cuda", dtype=torch.int32), + torch.tensor(starts, device="cuda", dtype=torch.int32), + torch.tensor(sum((list(range(count)) for count in counts), []), device="cuda", dtype=torch.int32), + ) + assert accept_lengths.tolist() == accepts + [1] + assert accepted_index.tolist() == expected_mask + for start, count, accepted in zip(starts[:2], counts[:2], accepts): + request = requests[start] + expected, states = reference_segment( + residual[start : start + count], + tokens[start : start + count], + source, + config, + before_ngram[request : request + 1], + before_conv[request : request + 1], + ) + torch.testing.assert_close(actual[start : start + count], expected, atol=0.016, rtol=0.016) + ngram, convolution = states[accepted] + torch.testing.assert_close(cache.ngram[request], ngram[0], atol=0, rtol=0) + torch.testing.assert_close(cache.convolution[request], convolution[0], atol=0.016, rtol=0.016) + assert req.req_to_mtp_state_index[request] == accepted - 1 + torch.testing.assert_close(cache.ngram[7], untouched.ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution[7], untouched.convolution, atol=0, rtol=0) + if capture: + graph.reset() + + +def test_microbatch_candidates_survive_reverse_verification(config): + cache, source = make_cache(config) + speculative = LexicalSpeculativeCache(cache, 3) + batches = [] + for request, ids in ((5, [19, 63, 23]), (2, [29, 31, 37])): + meta = metadata([request] * 3, [11, 12, 13], [1] * 3) + tokens = torch.tensor(ids, device="cuda") + residual = torch.randn(3, 64, device="cuda", dtype=torch.bfloat16) + batch = cache.prepare(meta) + batch.apply(residual, tokens, cache.layer_weight, 1, None) + speculative.stage(meta, batch) + batches.append((request, residual, tokens)) + for request, residual, tokens in reversed(batches): + _, states = reference_segment( + residual, + tokens, + source, + config, + cache.ngram[request : request + 1].clone(), + cache.convolution[request : request + 1].clone(), + ) + speculative.commit( + torch.tensor([request] * 3, device="cuda", dtype=torch.int32), + torch.tensor([0], device="cuda", dtype=torch.int32), + torch.tensor([2], device="cuda", dtype=torch.int32), + 7, + ) + torch.testing.assert_close(cache.ngram[request], states[2][0][0], atol=0, rtol=0) + torch.testing.assert_close(cache.convolution[request], states[2][1][0], atol=0.016, rtol=0.016) diff --git a/unit_tests/models/qwen4_exp/test_vision.py b/unit_tests/models/qwen4_exp/test_vision.py new file mode 100644 index 0000000000..ccf34a0521 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_vision.py @@ -0,0 +1,49 @@ +import pytest +import torch + +from lightllm.models.qwen4_exp.vision import Qwen4ExpVisionTransformer +from lightllm.server import visualserver +from lightllm.common.basemodel.attention_vit.create_utils import get_vit_att_backend_class +from .vision_reference import interpolate_positions, vision_reference + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_vision_encoder_matches_independent_packed_frames(dtype, monkeypatch): + monkeypatch.setattr(visualserver, "VIT_ATTN_BACKEND", get_vit_att_backend_class("sdpa")) + config = dict( + depth=2, + hidden_size=32, + out_hidden_size=16, + intermediate_size=48, + hidden_act="gelu_pytorch_tanh", + num_heads=4, + in_channels=3, + patch_size=2, + temporal_patch_size=2, + spatial_merge_size=2, + num_position_embeddings=16, + deepstack_visual_indexes=[], + ) + torch.manual_seed(1907) + model = Qwen4ExpVisionTransformer({"data_type": dtype}, **config).eval().to(device="cuda", dtype=dtype) + grid = torch.tensor([[1, 4, 6], [2, 2, 4]], device="cuda") + pixels = torch.randn(40, 24, device="cuda", dtype=dtype) + source = {name: value.detach().clone() for name, value in model.state_dict().items()} + tolerance = 0.016 if dtype == torch.bfloat16 else 2e-5 + with torch.no_grad(): + actual_positions = model.fast_pos_embed_interpolate(grid) + expected_positions = interpolate_positions(source["pos_embed.weight"], grid, 2) + torch.testing.assert_close(actual_positions, expected_positions, atol=tolerance, rtol=tolerance) + actual, deepstack = model(pixels, grid) + expected = vision_reference(pixels, grid, source, config) + torch.testing.assert_close(actual_positions, expected_positions, atol=tolerance, rtol=tolerance) + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + assert not deepstack + assert actual.shape == (10, 16) + + +def test_vision_rejects_unsupported_deepstack(): + with pytest.raises(ValueError, match="deepstack"): + Qwen4ExpVisionTransformer({"data_type": "bf16"}, deepstack_visual_indexes=[1]) diff --git a/unit_tests/models/qwen4_exp/test_weights.py b/unit_tests/models/qwen4_exp/test_weights.py new file mode 100644 index 0000000000..51bebb72ee --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_weights.py @@ -0,0 +1,99 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.qwen4_exp.checkpoint import normalize_text_weight_names, ngram_table_shape +from lightllm.models.qwen4_exp.layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight +from lightllm.models.qwen4_exp.layer_weights.pre_and_post_layer_weight import Qwen4ExpPreAndPostLayerWeight +from lightllm.models.qwen4_exp.layer_infer.pre_layer_infer import Qwen4ExpPreLayerInfer +from lightllm.models.qwen4_exp.layer_infer.post_layer_infer import Qwen4ExpPostLayerInfer +from lightllm.common.quantization import Quantcfg +from .helpers import checkpoint, prepost_checkpoint, reference_hyper + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def test_scalar_gate_dispatch_is_model_owned(config): + from lightllm.common.quantization.no_quant import NoQuantization + from lightllm.common.quantization.scalar_projection import StableScalarProjection + from lightllm.models.qwen3_5_moe.layer_weights.transformer_layer_weight import Qwen35MOETransformerLayerWeight + + native = Qwen4ExpTransformerLayerWeight(0, torch.bfloat16, config, Quantcfg(config)) + legacy = Qwen35MOETransformerLayerWeight(0, torch.bfloat16, config, Quantcfg(config)) + assert isinstance(native.shared_expert_gate.quant_method, StableScalarProjection) + assert type(legacy.shared_expert_gate.quant_method) is NoQuantization + + +@pytest.mark.parametrize("layer", [0, 1, 3]) +def test_incremental_native_checkpoint_loading(config, layer): + torch.manual_seed(737) + weights = checkpoint(config, layer) + native = Qwen4ExpTransformerLayerWeight(layer, torch.bfloat16, config, Quantcfg(config)) + assert not native.attn_hyper_connection.verify_load() + with pytest.raises(RuntimeError, match="not fully loaded"): + _ = native.attn_hyper_connection.component + + for key in reversed(list(weights)): + native.load_hf_weights({key: weights[key]}) + native.verify_load() + assert not hasattr(native, "att_norm_weight_") + assert not hasattr(native, "ffn_norm_weight_") + if layer == 3: + assert native.qk_norm_weight_.q_weight.dtype == torch.float32 + expected = 1 + weights[f"model.language_model.layers.{layer}.self_attn.q_norm.weight"].float().cuda() + torch.testing.assert_close(native.qk_norm_weight_.q_weight, expected, atol=0, rtol=0) + assert native.indexer.component.head_dim == 16 + else: + assert native.linear_norm.activation == "sigmoid" + hidden = torch.randn(3, 32, device="cuda", dtype=torch.bfloat16) + actual = native.linear_in_proj.mm(hidden, use_custom_tensor_mananger=False) + p = f"model.language_model.layers.{layer}.linear_attn" + expected = torch.cat( + [ + torch.nn.functional.linear(hidden, weights[p + f".in_proj_{name}.weight"].cuda()) + for name in ["qkv", "z", "b", "a"] + ], + -1, + ) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + if layer == 1: + assert native.ngram_embedding.weight.shape == (36, 8) + assert native.ngram_hasher.component.heads == 4 + assert native.ple.component.context_len == 9 + + +def test_embedding_widening_and_final_native_logits(config): + torch.manual_seed(737) + source = prepost_checkpoint(config) + native = Qwen4ExpPreAndPostLayerWeight(torch.bfloat16, config) + native.load_hf_weights(dict(source)) + native.verify_load() + assert not hasattr(native, "final_norm_weight_") + pre, post = Qwen4ExpPreLayerInfer(config), Qwen4ExpPostLayerInfer(config) + state = SimpleNamespace(is_prefill=False, is_draft_model=False, batch_size=2, prompt_logics=None, dist_group=None) + ids = torch.tensor([2, 17], device="cuda", dtype=torch.int64) + widened = pre.token_forward(ids, state, native) + torch.testing.assert_close( + widened, source["model.language_model.embed_tokens.weight"].cuda()[ids].repeat(1, 2), atol=0, rtol=0 + ) + widened = widened + torch.randn_like(widened) * 0.1 + actual = post.token_forward(widened, state, native).logits + mixed = reference_hyper(widened, source, "model.language_model.hyper_connection_mixer", config, False) + expected = torch.nn.functional.linear(mixed, source["lm_head.weight"].cuda()).float() + torch.testing.assert_close(actual, expected, atol=0.016, rtol=0.016) + + +def test_namespace_and_allocation_metadata(): + value = torch.zeros(1) + weights = { + "model.language_model.x": value, + "model.visual.x": value, + "mtp.x": value, + "adapter.model.language_model.x": value, + } + normalize_text_weight_names(weights) + assert set(weights) == {"model.x", "model.visual.x", "mtp.x", "adapter.model.language_model.x"} + with pytest.raises(ValueError, match="Conflicting"): + normalize_text_weight_names({"model.x": value, "model.language_model.x": value.clone()}) + assert ngram_table_shape(20000000, 16, 0, 2560, 128) == (320001536, 160) diff --git a/unit_tests/models/qwen4_exp/vision_reference.py b/unit_tests/models/qwen4_exp/vision_reference.py new file mode 100644 index 0000000000..bd0f2b68b1 --- /dev/null +++ b/unit_tests/models/qwen4_exp/vision_reference.py @@ -0,0 +1,93 @@ +import torch +import torch.nn.functional as F + + +def patch_coordinates(grid, merge, device): + coordinates = [] + for frames, height, width in grid.cpu().tolist(): + coordinates.extend( + (y + dy, x + dx) + for _ in range(frames) + for y in range(0, height, merge) + for x in range(0, width, merge) + for dy in range(merge) + for dx in range(merge) + ) + return torch.tensor(coordinates, dtype=torch.int64, device=device) + + +def interpolate_positions(table, grid, merge): + side = int(table.shape[0] ** 0.5) + result = [] + for frames, height, width in grid.cpu().tolist(): + coords = patch_coordinates(torch.tensor([[frames, height, width]]), merge, table.device) + ys = torch.linspace(0, side - 1, height, device=table.device)[coords[:, 0]] + xs = torch.linspace(0, side - 1, width, device=table.device)[coords[:, 1]] + y0, x0 = ys.floor().long(), xs.floor().long() + y1, x1 = (y0 + 1).clamp_max(side - 1), (x0 + 1).clamp_max(side - 1) + dy, dx = ys - y0, xs - x0 + indices = torch.stack([y0 * side + x0, y0 * side + x1, y1 * side + x0, y1 * side + x1], -1) + coefficients = torch.stack([(1 - dy) * (1 - dx), (1 - dy) * dx, dy * (1 - dx), dy * dx], -1) + result.append((table[indices] * coefficients[:, :, None]).sum(1).to(table.dtype)) + return torch.cat(result) + + +def vision_reference(pixels, grid, weights, config): + hidden, heads = config["hidden_size"], config["num_heads"] + patch, temporal, merge = config["patch_size"], config["temporal_patch_size"], config["spatial_merge_size"] + values = F.conv3d( + pixels.view(-1, config["in_channels"], temporal, patch, patch), + weights["patch_embed.proj.weight"], + weights["patch_embed.proj.bias"], + stride=(temporal, patch, patch), + ).view(-1, hidden) + values = values + interpolate_positions(weights["pos_embed.weight"], grid, merge) + dim = hidden // heads + theta = config.get("rope_parameters", {}).get("rope_theta", 10000.0) + invfreq = theta ** (-torch.arange(0, dim // 2, 2, device=pixels.device).float() / (dim // 2)) + coordinates = patch_coordinates(grid, merge, pixels.device) + angles = (coordinates[:, :, None] * invfreq).flatten(1).repeat(1, 2) + cos, sin = angles.cos()[:, None], angles.sin()[:, None] + + def norm(value, prefix): + return F.layer_norm(value, (value.shape[-1],), weights[prefix + ".weight"], weights[prefix + ".bias"], 1e-6) + + def rotate(value): + value = value.float() + a, b = value.chunk(2, -1) + return (value * cos + torch.cat([-b, a], -1) * sin).to(pixels.dtype) + + lengths = [height * width for frames, height, width in grid.cpu().tolist() for _ in range(frames)] + for layer in range(config["depth"]): + prefix = f"blocks.{layer}" + qkv = F.linear( + norm(values, prefix + ".norm1"), weights[prefix + ".attn.qkv.weight"], weights[prefix + ".attn.qkv.bias"] + ) + q, k, v = qkv.view(-1, 3, heads, dim).unbind(1) + q, k = rotate(q), rotate(k) + outputs, offset = [], 0 + for count in lengths: + outputs.append( + F.scaled_dot_product_attention( + q[offset : offset + count].transpose(0, 1)[None], + k[offset : offset + count].transpose(0, 1)[None], + v[offset : offset + count].transpose(0, 1)[None], + )[0].transpose(0, 1) + ) + offset += count + attention = torch.cat(outputs).reshape(-1, hidden) + values = values + F.linear( + attention, weights[prefix + ".attn.proj.weight"], weights[prefix + ".attn.proj.bias"] + ) + mlp = F.linear( + norm(values, prefix + ".norm2"), + weights[prefix + ".mlp.linear_fc1.weight"], + weights[prefix + ".mlp.linear_fc1.bias"], + ) + mlp = F.gelu(mlp, approximate="tanh" if config["hidden_act"] == "gelu_pytorch_tanh" else "none") + values = values + F.linear( + mlp, weights[prefix + ".mlp.linear_fc2.weight"], weights[prefix + ".mlp.linear_fc2.bias"] + ) + values = norm(values, "merger.norm").view(-1, hidden * merge * merge) + values = F.gelu(F.linear(values, weights["merger.linear_fc1.weight"], weights["merger.linear_fc1.bias"])) + return F.linear(values, weights["merger.linear_fc2.weight"], weights["merger.linear_fc2.bias"]) diff --git a/unit_tests/models/sensenova_loop/test_gdn_state.py b/unit_tests/models/sensenova_loop/test_gdn_state.py new file mode 100644 index 0000000000..a31c0acaf8 --- /dev/null +++ b/unit_tests/models/sensenova_loop/test_gdn_state.py @@ -0,0 +1,150 @@ +from types import 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.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 + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("mtp_step", [0, 2]) +@pytest.mark.parametrize( + "start,end,count,logical_layer,storage_layer", + [(4, 28, 1, 4, 3), (4, 28, 3, 28, 21), (1, 6, 2, 13, 10)], + ids=["no-loop", "aligned-loop", "unaligned-loop"], +) +def test_gdn_decode_keeps_logical_layer_state( + monkeypatch, device, mtp_step, start, end, count, logical_layer, storage_layer +): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + torch.manual_seed(1585) + settings = dict(start=start, end=end, eval_count=count, use_gate=False) + config = dict( + num_hidden_layers=32, + full_attention_interval=4, + layer_types=["full_attention" if (i + 1) % 4 == 0 else "linear_attention" for i in range(32)], + linear_num_key_heads=1, + linear_num_value_heads=2, + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_conv_kernel_dim=4, + hidden_act="silu", + ) + indices = get_loop_layer_indices(settings, 32) + kinds = get_loop_layer_types(settings, config) + layers = tuple(i for i, kind in enumerate(kinds) if kind == "linear_attention") + args = StartArgs(mtp_step=mtp_step, linear_att_ssm_data_type="float32") + monkeypatch.setattr("lightllm.common.basemodel.attention.linear.gdn.get_env_start_args", lambda: args) + monkeypatch.setattr("lightllm.common.basemodel.attention.base_att.get_env_start_args", lambda: args) + backend = object.__new__(TritonLinearAttBackend) + backend.__init__( + SimpleNamespace( + config=config, + tp_world_size_=1, + is_mtp_draft_model=False, + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: mtp_step), + ) + ) + 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) + resources = AttentionStateManager(request_capacity=3) + resources.register_state("main", resource) + resources.finish_state_registration() + assert resource.get_storage_layer_index(logical_layer) == storage_layer + physical_weights = [ + SimpleNamespace( + layer_num_=i, + linear_conv1d=SimpleNamespace( + mm_param=SimpleNamespace(weight=torch.randn(layout.conv_channels, 4, device=device, dtype=dtype) * 0.1), + bias=None, + ), + linear_A_log=SimpleNamespace(weight=torch.zeros(2, device=device)), + linear_dt_bias=SimpleNamespace(weight=torch.zeros(2, device=device)), + ) + for i in range(32) + ] + execution_weights = [physical_weights[i] for i in indices] + weight = execution_weights[logical_layer] + assert weight is physical_weights[indices[logical_layer]] + + if device == "cpu": + + def conv_pack(mixed, z, a, b, conv, *args): + conv[1].add_(1) + return (*backend._rearrange_mixed_qkv(mixed, decode=True), z, a, b) + + def conv_mtp(mixed, conv, *args, **kwargs): + conv[1].add_(1) + return mixed + + 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 + ) + monkeypatch.setattr( + "lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent_gated_delta_rule", recurrent + ) + monkeypatch.setattr( + "lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent." + "mtp_fused_recurrent_gated_delta_rule", + recurrent, + ) + + infer = InferStateInfo() + infer.att_state_manager, infer.att_state_role = resources, "main" + infer.b_req_idx = torch.tensor([1], device=device, dtype=torch.int32) + control = dict(layer_num=logical_layer, layer_weight=weight) + 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) + ), + ), + ) + + 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) + 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) + ), + ), + ) + assert torch.isfinite(output).all() + assert not torch.equal(resource.conv_state[storage_layer], before_conv[storage_layer]) + assert not torch.equal(resource.ssm_state[storage_layer], before_ssm[storage_layer]) + assert torch.equal(resource.conv_state[untouched], before_conv[untouched]) + assert torch.equal(resource.ssm_state[untouched], before_ssm[untouched]) + assert weight.layer_num_ == indices[logical_layer] diff --git a/unit_tests/server/httpserver/test_slot_availability.py b/unit_tests/server/httpserver/test_slot_availability.py new file mode 100644 index 0000000000..74b687a72e --- /dev/null +++ b/unit_tests/server/httpserver/test_slot_availability.py @@ -0,0 +1,274 @@ +import asyncio +from types import SimpleNamespace +import pytest +import torch +from lightllm.server.httpserver.slot_availability import RequestSlotAvailability +from lightllm.server.core.objs.shm_req_manager import ShmReqManager + +from lightllm.server.httpserver.manager import HttpServerManager +import lightllm.server.httpserver.manager as module + + +class Pool(ShmReqManager): + def __init__(self, indices=()): + self.free = list(indices) + self.claimed = set() + self.empty = asyncio.Event() + self.fail_release = False + + def alloc_req_index(self): + if not self.free: + self.empty.set() + return None + index = self.free.pop(0) + assert index not in self.claimed + self.claimed.add(index) + return index + + def alloc_req_indexes(self, count): + if len(self.free) < count: + self.empty.set() + return None + return [self.alloc_req_index() for _ in range(count)] + + def release_req_index(self, index): + if self.fail_release: + raise RuntimeError("test release failed") + self.claimed.discard(index) + self.free.append(index) + + async def async_put_back_req_obj(self, req): + pass + + def __del__(self): + pass + + +def manager(pool, availability=True): + result = HttpServerManager.__new__(HttpServerManager) + result.shm_req_manager = pool + result._request_slot_availability = RequestSlotAvailability(availability) + result.is_multinode_tp_slave = False + result.args = SimpleNamespace(running_max_req_size=8, run_mode="normal") + return result + + +@pytest.fixture(autouse=True) +def no_cuda(): + assert not torch.cuda.is_initialized() + yield + assert not torch.cuda.is_initialized() + + +def test_original_async_allocator_keeps_immediate_none_contract(): + async def scenario(): + assert await Pool().async_alloc_req_index() is None + + asyncio.run(scenario()) + + +def test_generation_change_before_wait_is_observed(): + async def scenario(): + notifier = RequestSlotAvailability() + generation = notifier.generation + notifier.notify_release() + assert await notifier.wait(generation, 10) + + asyncio.run(scenario()) + + +def test_waiting_task_is_notified_after_release(): + async def scenario(): + notifier = RequestSlotAvailability() + task = asyncio.create_task(notifier.wait(0, 10)) + await asyncio.sleep(0) + notifier.notify_release() + assert await task + + asyncio.run(scenario()) + + +def test_unobserved_external_release_uses_timeout_fallback(): + async def scenario(): + notifier = RequestSlotAvailability() + assert not await notifier.wait(notifier.generation, 0.001) + + asyncio.run(scenario()) + + +def test_actual_http_retry_observes_release_between_check_and_wait(): + async def scenario(): + pool = Pool() + http = manager(pool) + availability = http._request_slot_availability + original = availability.wait + + async def release_before_wait(generation, timeout): + await pool.async_release_req_index(7) + http._notify_req_slot_release() + return await original(generation, timeout) + + availability.wait = release_before_wait + assert await http._alloc_shm_req_indexes(1) == [7] + assert pool.claimed == {7} + + asyncio.run(scenario()) + + +def test_recycler_notifies_only_after_actual_release(): + async def scenario(): + pool = Pool() + pool.claimed = {7} + http = manager(pool) + http.recycle_event = asyncio.Event() + req = SimpleNamespace(request_id=8, index_in_shm_mem=7) + group = SimpleNamespace(group_req_id=8, shm_req_objs=[req], multimodal_params=None) + http.req_id_to_out_inf = {8: SimpleNamespace(can_release=lambda: True, group_req_objs=group)} + + async def release_media(params): + pass + + http._release_multimodal_resources = release_media + waiter = asyncio.create_task(http._alloc_shm_req_indexes(1)) + await pool.empty.wait() + recycler = asyncio.create_task(http.recycle_resource_loop()) + http.recycle_event.set() + try: + assert await waiter == [7] + assert not http.req_id_to_out_inf + assert http._request_slot_availability.generation == 1 + finally: + recycler.cancel() + with pytest.raises(asyncio.CancelledError): + await recycler + + asyncio.run(scenario()) + + +def test_http_cancellation_keeps_unallocated_group_free(): + async def scenario(): + pool = Pool([3]) + http = manager(pool) + task = asyncio.create_task(http._alloc_shm_req_indexes(2)) + await pool.empty.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert pool.free == [3] and not pool.claimed + assert http._request_slot_availability.generation == 0 + + asyncio.run(scenario()) + + +def test_waiter_cancellation_does_not_consume_notification(): + async def scenario(): + notifier = RequestSlotAvailability() + task = asyncio.create_task(notifier.wait(0, 10)) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + notifier.notify_release() + assert await notifier.wait(0, 10) + + asyncio.run(scenario()) + + +def test_multiple_http_waiters_do_not_duplicate_allocation(): + async def scenario(): + pool = Pool() + http = manager(pool) + tasks = [asyncio.create_task(http._alloc_shm_req_indexes(1)) for _ in range(2)] + await pool.empty.wait() + await pool.async_release_req_index(4) + http._notify_req_slot_release() + await pool.async_release_req_index(5) + http._notify_req_slot_release() + assert {result[0] for result in await asyncio.gather(*tasks)} == {4, 5} + assert pool.claimed == {4, 5} + + asyncio.run(scenario()) + + +def test_failed_release_does_not_increment_generation(): + async def scenario(): + pool = Pool([3]) + http = manager(pool) + pool.fail_release = True + + task = asyncio.create_task(http._alloc_shm_req_indexes(2)) + await pool.empty.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert http._request_slot_availability.generation == 0 + + asyncio.run(scenario()) + + +@pytest.mark.parametrize( + "change", + [ + dict(run_mode="prefill"), + dict(run_mode="decode"), + dict(dp=2), + dict(nnodes=2), + dict(enable_rl=True), + dict(httpserver_workers=2), + ], +) +def test_unverified_modes_keep_legacy_backoff(change): + args = SimpleNamespace(run_mode="normal", dp=1, nnodes=1, enable_rl=False, httpserver_workers=1) + for key, value in change.items(): + setattr(args, key, value) + notifier = RequestSlotAvailability.from_args(args) + assert not notifier.enabled + notifier.notify_release() + assert notifier.generation == 0 + + +@pytest.mark.parametrize("priority,expected", [(False, 0.1), (True, 0.02)]) +def test_original_priority_backoff_interval_is_preserved(priority, expected): + async def scenario(): + pool = Pool() + http = manager(pool, False) + observed = [] + + async def wait_and_release(generation, timeout): + observed.append(timeout) + await pool.async_release_req_index(7) + return False + + http._request_slot_availability.wait = wait_and_release + assert await http._alloc_shm_req_indexes(1, high_priority_request=priority) == [7] + assert observed == pytest.approx([expected]) + + asyncio.run(scenario()) + + +def test_existing_timeout_preserves_unallocated_group(): + async def scenario(): + pool = Pool([3]) + http = manager(pool, False) + http.pd_node_request_limit_enabled = True + http.pd_node_shm_req_alloc_timeout_seconds = 0 + http.args = SimpleNamespace(run_mode="prefill", running_max_req_size=8) + with pytest.raises(module.ServerBusyError): + await http._alloc_shm_req_indexes(2, pd_node_resource_wait_timeout_seconds=0) + assert pool.free == [3] and not pool.claimed + + asyncio.run(scenario()) + + +@pytest.mark.parametrize("mode,workers,enabled", [("normal", 1, True), ("decode", 1, False), ("normal", 2, False)]) +def test_real_constructor_initializes_notifier_before_external_io(monkeypatch, mode, workers, enabled): + args = SimpleNamespace(run_mode=mode, dp=1, nnodes=1, enable_rl=False, httpserver_workers=workers) + + def stop_before_external_io(): + raise RuntimeError("stop_before_external_io") + + monkeypatch.setattr(module, "get_shm_port_args", stop_before_external_io) + http = HttpServerManager.__new__(HttpServerManager) + with pytest.raises(RuntimeError, match="stop_before_external_io"): + http.__init__(args) + assert http._request_slot_availability.enabled is enabled diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py b/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py index 2e488ebb30..9a64edf769 100644 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py @@ -962,3 +962,86 @@ def test_topk_prefix_sums_only_computes_requested_counts(): assert set(result) == {0, 2, 4} np.testing.assert_allclose([result[0], result[2], result[4]], [0.0, 1.6, 2.1]) + + +def build_stalled_zero_draft_planner(req_num, spec_mode="eagle_no_att"): + planner = build_lightspec_planner(max_draft_step=4, spec_mode=spec_mode) + for batch_size in range(1, 41): + planner.target_infer_costs.update(batch_size, 1.0 + 0.1 * batch_size) + planner.draft_infer_costs.update(batch_size, 0.2 + 0.01 * batch_size) + planner._update_verified_batch( + accept_lengths=[1] * req_num, + req_num=req_num, + dynamic_batch_size=5 * req_num, + verified_draft_step=4, + ) + assert planner._select_draft_step(req_num) == planner.draft_steps[0] + planner.pre_draft_step = planner.draft_steps[0] + return planner + + +@pytest.mark.parametrize("req_num", [1, 2, 8]) +@pytest.mark.parametrize("all_new", [False, True]) +def test_lightspec_new_request_recovers_from_zero_draft_feedback(req_num, all_new): + planner = build_stalled_zero_draft_planner(req_num) + target_costs = dict(planner.target_infer_costs.infer_cost_ms_table) + draft_costs = dict(planner.draft_infer_costs.infer_cost_ms_table) + new_requests = req_num if all_new else 1 + + plan = planner.plan( + decode_reqs=build_decode_reqs(req_num, req_num_with_proposals=req_num - new_requests), + origin_batch_size=5 * req_num, + ) + + assert plan.pre_draft_step == 0 + assert plan.dynamic_batch_size == req_num + assert plan.draft_step == 4 + assert not plan.all_reqs_have_proposals + assert dict(planner.target_infer_costs.infer_cost_ms_table) == target_costs + assert dict(planner.draft_infer_costs.infer_cost_ms_table) == draft_costs + + planner.pre_draft_step = plan.draft_step + next_plan = planner.plan(decode_reqs=build_decode_reqs(req_num), origin_batch_size=5 * req_num) + assert next_plan.dynamic_batch_size == 5 * req_num + planner._update_verified_batch( + accept_lengths=[5] * req_num, + req_num=req_num, + dynamic_batch_size=next_plan.dynamic_batch_size, + verified_draft_step=4, + ) + assert planner._select_draft_step(req_num) == 4 + + +@pytest.mark.parametrize("empty", [False, True]) +def test_lightspec_existing_or_empty_batch_retains_zero_draft_evidence(empty): + planner = build_stalled_zero_draft_planner(2) + progress = dict(planner.progress_ema_by_config) + survival = list(planner.prefix_survival_by_depth) + + plan = planner.plan(decode_reqs=[] if empty else build_decode_reqs(2), origin_batch_size=0 if empty else 10) + + assert plan.draft_step == 0 + assert planner.progress_ema_by_config == progress + assert planner.prefix_survival_by_depth == survival + + +@pytest.mark.parametrize("req_num", [1, 2, 8]) +@pytest.mark.parametrize("spec_mode", ["eagle_with_att", "eagle3", "vanilla_with_att"]) +def test_lightspec_new_workload_escapes_target_only_verification(req_num, spec_mode): + planner = build_stalled_zero_draft_planner(req_num, spec_mode=spec_mode) + previous_depth = planner.pre_draft_step + assert previous_depth > 0 + target_costs = dict(planner.target_infer_costs.infer_cost_ms_table) + draft_costs = dict(planner.draft_infer_costs.infer_cost_ms_table) + + first = planner.plan(decode_reqs=build_decode_reqs(req_num, 0), origin_batch_size=5 * req_num) + + assert first.pre_draft_step == previous_depth + assert first.dynamic_batch_size == req_num + assert first.draft_step == 4 + assert not first.all_reqs_have_proposals + assert dict(planner.target_infer_costs.infer_cost_ms_table) == target_costs + assert dict(planner.draft_infer_costs.infer_cost_ms_table) == draft_costs + planner.pre_draft_step = first.draft_step + following = planner.plan(decode_reqs=build_decode_reqs(req_num), origin_batch_size=5 * req_num) + assert following.dynamic_batch_size == 5 * req_num diff --git a/unit_tests/server/router/req_queue/test_cached_prefix_queue.py b/unit_tests/server/router/req_queue/test_cached_prefix_queue.py new file mode 100644 index 0000000000..b941f91475 --- /dev/null +++ b/unit_tests/server/router/req_queue/test_cached_prefix_queue.py @@ -0,0 +1,224 @@ +import random +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +import lightllm.server.router.req_queue as module +from lightllm.server.router.req_queue.chunked_prefill.cached_prefix_queue import ( + CachedPrefixChunkedPrefillQueue, + supports_cached_prefix_reservations, +) +from lightllm.server.router.req_queue.chunked_prefill.impl import ChunkedPrefillQueue +from lightllm.server.router.req_queue.cached_prefix import CachedPrefixReservations + + +class Load: + def set_estimated_peak_token_count(self, value, index): + self.peak = value + + def set_dynamic_max_load(self, value, index): + self.ratio = value + + +class Request: + def __init__(self, identity, base=1800, remaining=350, cached=512): + self.request_id, self.input_len, self.shm_cur_kv_len = identity, base, base + self.prompt_cache_len = cached + self.is_paused = self.is_aborted = False + self.sample_params = SimpleNamespace(suggested_dp_index=0, disable_prompt_cache=False) + self.multimodal_params = SimpleNamespace(images=[], audios=[]) + self.shm_prompt_ids = SimpleNamespace(arr=np.arange(base, dtype=np.int64)) + self.linear_att_token_hash_list = SimpleNamespace(get_all=lambda: [17] * (base // 512)) + self.remaining = remaining + + def get_tuple_tokens(self, is_busy, ema): + return self.input_len, self.remaining + + def get_first_router_need_tokens(self): + return 512 + + +def arguments(**changes): + result = SimpleNamespace( + dp=1, + nnodes=1, + run_mode="normal", + enable_rl=False, + disable_dynamic_prompt_cache=False, + enable_cpu_cache=False, + linear_att_page_block_num=1, + linear_att_hash_page_size=512, + short_prefill_token_threshold=None, + diverse_mode=False, + output_constraint_mode="none", + first_token_constraint_mode=False, + disable_chunked_prefill=False, + model_dir="unused", + ) + for name, value in changes.items(): + setattr(result, name, value) + return result + + +def queue(requests, cls=CachedPrefixChunkedPrefillQueue): + result = cls.__new__(cls) + result.args = arguments() + result.dp_index = 0 + result.max_total_tokens, result.running_max_req_size, result.batch_max_tokens = 16384, 8, 1024 + result.router = SimpleNamespace(router_statics=SimpleNamespace(ema_req_out_len=2048), shared_token_load=Load()) + result.is_busy = lambda: True + if cls is CachedPrefixChunkedPrefillQueue: + result.prefix_reservations = CachedPrefixReservations(result.args, 0) + result._init_cache_list(SimpleNamespace(reqs=requests), True) + return result + + +@pytest.fixture(autouse=True) +def namespace(monkeypatch): + assert not torch.cuda.is_initialized() + monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", "owned-canonical-test") + yield + assert not torch.cuda.is_initialized() + + +def test_known_cached_prefix_changes_actual_admission_at_same_capacity(): + legacy = queue([Request(i) for i in range(7)], ChunkedPrefillQueue) + candidate = queue([Request(i) for i in range(7)]) + assert not legacy._can_add_new_req(Request(7), True) + assert candidate._can_add_new_req(Request(7), True) is True + assert candidate.router.shared_token_load.peak == 13616 + assert candidate.max_total_tokens == 16384 and candidate.running_max_req_size == 8 + + +def test_refactored_legacy_peak_matches_original_formula_for_ties_and_empty_batches(): + rng = random.Random(8245) + for _ in range(1500): + entries = [(rng.randrange(1, 3000), rng.randrange(0, 100)) for _ in range(rng.randrange(0, 17))] + legacy = queue([], ChunkedPrefillQueue) + legacy.cache_len_list = entries.copy() + expected = max( + ( + sum(a for a, b in sorted(entries, key=lambda e: -e[1])[:i]) + i * pair[1] + for i, pair in enumerate(sorted(entries, key=lambda e: -e[1]), 1) + ), + default=0, + ) + assert legacy._estimate_peak_tokens() == expected + + +@pytest.mark.parametrize( + "change", + [ + "images", + "audios", + "paused", + "aborted", + "disabled_cache", + "unresident", + "missing_hashes", + "different_hashes", + "different_tokens", + ], +) +def test_uncertain_or_distinct_prefixes_preserve_legacy_rejection(change): + requests = [Request(i) for i in range(8)] + for request in requests: + if change in ("images", "audios"): + setattr(request.multimodal_params, change, [object()]) + elif change in ("paused", "aborted"): + setattr(request, "is_" + change, True) + elif change == "disabled_cache": + request.sample_params.disable_prompt_cache = True + elif change == "unresident": + request.shm_cur_kv_len = 128 + elif change == "missing_hashes": + request.linear_att_token_hash_list = SimpleNamespace(get_all=lambda: []) + elif change == "different_hashes": + request.linear_att_token_hash_list = SimpleNamespace(get_all=lambda i=request.request_id: [i]) + elif change == "different_tokens": + request.shm_prompt_ids.arr += request.request_id * 1800 + candidate = queue(requests[:7]) + assert not candidate._can_add_new_req(requests[7], True) + + +def test_reset_and_load_preserve_tuple_identity_and_request_lifetime(): + requests = [Request(i, remaining=350 - i * 3) for i in range(7)] + candidate = queue(requests) + assert candidate._can_add_new_req(Request(7, remaining=420), True) + peak, ratio = candidate._calcu_batch_token_load_batch_not_none(SimpleNamespace(reqs=requests)) + assert peak < 16384 and ratio == peak / 16384 + candidate._init_cache_list(SimpleNamespace(reqs=[Request(50)]), True) + assert set(candidate.prefix_reservations.memo) == {50} + assert len(candidate.prefix_reservations.by_tuple) == 1 + candidate._init_cache_list(None, True) + assert not candidate.prefix_reservations.memo and candidate._estimate_peak_tokens() == 0 + + +@pytest.mark.parametrize( + "change", + [ + dict(dp=2), + dict(nnodes=2), + dict(run_mode="decode"), + dict(enable_rl=True), + dict(disable_dynamic_prompt_cache=True), + dict(enable_cpu_cache=True), + dict(linear_att_page_block_num=4), + dict(linear_att_hash_page_size=0), + ], +) +def test_unverified_cache_geometry_and_modes_are_not_enabled(change): + assert not supports_cached_prefix_reservations(arguments(**change)) + + +@pytest.mark.parametrize( + "model_type,expected", + [("qwen4_exp", True), ("qwen4_exp_text", True), ("qwen3_5", False), ("llama", False), (None, False)], +) +def test_factory_selects_only_verified_models(monkeypatch, model_type, expected): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda path: model_type) + assert (module._get_req_queue_class(arguments(), None, 1) is CachedPrefixChunkedPrefillQueue) == expected + + +@pytest.mark.parametrize( + "change", + [ + dict(linear_att_page_block_num=4), + dict(enable_cpu_cache=True), + dict(dp=2), + dict(output_constraint_mode="x"), + dict(first_token_constraint_mode=True), + ], +) +def test_factory_retains_legacy_queue_for_unsupported_configurations(monkeypatch, change): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda path: "qwen4_exp") + assert module._get_req_queue_class(arguments(**change), None, 1) is ChunkedPrefillQueue + + +def test_request_limit_is_unchanged(): + candidate = queue([Request(i) for i in range(8)]) + assert not candidate._can_add_new_req(Request(8), True) + + +def test_unknown_namespace_gets_no_credit(monkeypatch): + monkeypatch.delenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID") + candidate = queue([Request(i) for i in range(7)]) + assert not candidate._can_add_new_req(Request(7), True) + + +def test_missing_prompt_handle_attaches_via_existing_request_api(): + requests = [Request(i) for i in range(8)] + linked = [] + for request in requests: + del request.shm_prompt_ids + + def attach(request=request): + linked.append(request.request_id) + request.shm_prompt_ids = SimpleNamespace(arr=np.arange(request.input_len, dtype=np.int64)) + + request.link_prompt_ids_shm_array = attach + candidate = queue(requests[:7]) + assert candidate._can_add_new_req(requests[7], True) + assert linked == list(range(8)) diff --git a/unit_tests/server/router/req_queue/test_peak_reservation.py b/unit_tests/server/router/req_queue/test_peak_reservation.py new file mode 100644 index 0000000000..713b30a44e --- /dev/null +++ b/unit_tests/server/router/req_queue/test_peak_reservation.py @@ -0,0 +1,96 @@ +import random + +import pytest + +from lightllm.server.router.req_queue.peak_reservation import Reservation, estimate_peak + + +def simulated_peak(entries, page_tokens): + totals = [] + for progress in range(max((entry.remaining_tokens for entry in entries), default=0) + 1): + pages, private = set(), 0 + for entry in entries: + if entry.remaining_tokens < progress: + continue + private += entry.base_tokens + progress + if entry.cache_domain is not None: + private -= len(entry.prefix_pages) * page_tokens + for depth in range(1, len(entry.prefix_pages) + 1): + pages.add((entry.cache_domain, entry.prefix_pages[:depth])) + totals.append(private + len(pages) * page_tokens) + return max(totals) + + +def test_shared_prefix_reservation_changes_admission_without_increasing_capacity(): + entries = [Reservation(1800, 350, (b"same exact page",), b"one cache") for _ in range(8)] + peak = estimate_peak(entries, 512) + assert peak.legacy_tokens == 17200 + assert peak.shared_tokens == simulated_peak(entries, 512) == 13616 + assert peak.shared_tokens < 16384 < peak.legacy_tokens + + +def test_prefix_ancestry_and_cache_domain_prevent_false_sharing(): + entries = [ + Reservation(20, 3, (b"a", b"same"), b"cache"), + Reservation(20, 3, (b"b", b"same"), b"cache"), + Reservation(20, 3, (b"a", b"same"), b"other cache"), + ] + peak = estimate_peak(entries, 4) + assert peak.legacy_tokens == peak.shared_tokens == simulated_peak(entries, 4) + + +def test_unknown_domains_keep_legacy_reservation(): + entries = [Reservation(20, remaining, (b"same",)) for remaining in (1, 3, 7)] + peak = estimate_peak(entries, 4) + assert peak.legacy_tokens == peak.shared_tokens == simulated_peak(entries, 4) + + +def test_credits_follow_each_live_subset_when_requests_complete(): + entries = [ + Reservation(25, 2, (b"a", b"b", b"c"), b"cache"), + Reservation(25, 8, (b"a",), b"cache"), + Reservation(25, 5, (b"a", b"b"), b"cache"), + ] + assert estimate_peak(entries, 4).shared_tokens == simulated_peak(entries, 4) + + +def test_reservation_matches_independent_simulated_lifetimes(): + rng = random.Random(381020) + for _ in range(1500): + page_tokens = rng.choice((1, 2, 4, 512)) + entries = [] + for _ in range(rng.randrange(1, 9)): + pages = tuple(rng.choice((b"a", b"b", b"c")) for _ in range(rng.randrange(5))) + entries.append( + Reservation( + len(pages) * page_tokens + rng.randrange(8), + rng.randrange(13), + pages, + rng.choice((None, b"x", b"y")), + ) + ) + peak = estimate_peak(entries, page_tokens) + assert peak.shared_tokens == simulated_peak(entries, page_tokens) + assert 0 <= peak.shared_tokens <= peak.legacy_tokens + + +@pytest.mark.parametrize( + "entry,page_tokens", + [ + (Reservation(-1, 1), 1), + (Reservation(1, -1), 1), + (Reservation(True, 1), 1), + (Reservation(1, 1), 0), + (Reservation(1, 1), True), + (Reservation(1, 1, (b"a", b"b"), b"cache"), 1), + (Reservation(1, 1, ("text is not an exact byte identity",), b"cache"), 1), + (Reservation(1, 1, (b"a",), "mutable or ambiguous domain"), 1), + ], +) +def test_rejects_invalid_accounting_metadata(entry, page_tokens): + with pytest.raises(ValueError): + estimate_peak([entry], page_tokens) + + +def test_empty_batch_has_no_reservation(): + assert estimate_peak([], 512).shared_tokens == 0 diff --git a/unit_tests/server/test_hybrid_start_args.py b/unit_tests/server/test_hybrid_start_args.py index 230e23b0e5..1c391c09cd 100644 --- a/unit_tests/server/test_hybrid_start_args.py +++ b/unit_tests/server/test_hybrid_start_args.py @@ -102,3 +102,37 @@ def test_whole_image_blocks_fit_batch_at_startup( else: with pytest.raises(AssertionError, match="batch_max_tokens >= max_image_token_count"): _launch_subprocesses(args) + + +@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) + model_dir = "s3://test-bucket/qwen" + config = {"model_type": "qwen3_5_text"} + downloads = [] + + def list_files(path, extension): + assert path == model_dir + return ["config.json"] if extension == "json" else [] + + def download_json(remote, local): + downloads.append(remote) + Path(local).write_text(json.dumps(config)) + + monkeypatch.setattr(PetrelHelper, "list", list_files) + monkeypatch.setattr(PetrelHelper, "download_json", download_json) + args = launch_args("qwen3_5_text", ssm_state_mode=mode, mtp_step=step, mtp_mode="eagle_with_att" if step else None) + args.model_dir = model_dir + assert not Path(model_dir, "config.json").exists() + with pytest.raises(RuntimeError, match="startup validation passed"): + _launch_subprocesses(args) + assert downloads == [model_dir + "/config.json"] + if mode == "replay": + expected_compute_mode = "precompute" if step > 0 else "inline" + assert args.ssm_replay_config == { + "cache_len": 16, + "compute_mode": expected_compute_mode, + } diff --git a/unit_tests/server/test_native_stream_details.py b/unit_tests/server/test_native_stream_details.py new file mode 100644 index 0000000000..0d983258a0 --- /dev/null +++ b/unit_tests/server/test_native_stream_details.py @@ -0,0 +1,69 @@ +import asyncio +import copy +import json +from types import SimpleNamespace + +import pytest + +from lightllm.server import api_lightllm + + +@pytest.mark.parametrize("details", [None, False, True]) +@pytest.mark.parametrize("reported", [False, True]) +def test_stream_details_are_opt_in_and_preserve_metadata(monkeypatch, details, reported): + class Sampling: + best_of = 1 + + def init(self, **kwargs): + assert "return_details" not in kwargs + + async def verify_cpu_async(self, compiler): + pass + + monkeypatch.setattr(api_lightllm, "SamplingParams", Sampling) + parameters = {"do_sample": False} + if details is not None: + parameters["return_details"] = details + + class Request: + async def json(self): + return {"inputs": [11, 12], "parameters": parameters} + + metadata = dict( + id=7, + prompt_tokens=2, + count_output_tokens=1, + logprob=-0.2, + logprobs={}, + prompt_cache_len=1, + mtp_accepted_token_num=2, + mtp_verify_token_num=3, + mtp_verify_step_num=1, + input_usage=dict(input_text_tokens=2), + ) + fields = ("prompt_cache_len", "mtp_accepted_token_num", "mtp_verify_token_num", "mtp_verify_step_num") + if not reported: + for name in fields: + metadata.pop(name) + before = copy.deepcopy(metadata) + + async def generate(*args, **kwargs): + yield 0, "token", metadata, SimpleNamespace(is_finished=lambda: True, get_finish_reason=lambda: "length") + + async def collect(): + response = await api_lightllm.lightllm_generate_stream( + Request(), SimpleNamespace(tokenizer=None, generate=generate, output_grammar_compiler=None) + ) + return [json.loads(chunk.decode().removeprefix("data:")) async for chunk in response.body_iterator] + + bodies = asyncio.run(collect()) + assert len(bodies) == 1 + body = bodies[0] + for name in fields: + if details and reported: + assert body["token"][name] == metadata[name] + else: + assert name not in body["token"] + assert body["input_usage"] == metadata["input_usage"] and body["details"] is None + assert body["token"]["id"] == 7 and body["finish_reason"] == "length" + assert metadata == before diff --git a/unit_tests/server/test_ssm_state_mode.py b/unit_tests/server/test_ssm_state_mode.py new file mode 100644 index 0000000000..da786bddef --- /dev/null +++ b/unit_tests/server/test_ssm_state_mode.py @@ -0,0 +1,167 @@ +import json + +import pytest + +from lightllm.server.api_cli import make_argument_parser +from lightllm.server.api_start import _validate_and_normalize_ssm_config +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.mark.parametrize( + "mode,dtype,step,model,capacity,error", + [ + ("native", "bfloat16", 0, "llama", 16, None), + ("native", "bfloat16", 32, "llama", 16, None), + ("replay", "float32", 0, "sensenova_loop", 16, None), + ("replay", "bfloat16", 2, "sensenova_loop", 8, None), + ("replay", "float32", 0, "glm5_next", 8, None), + ("replay", "bfloat16", 2, "glm5_next_text", 8, None), + ("replay", "float32", 0, "qwen3_next", 16, None), + ("replay", "float32", 15, "qwen3_5", 16, None), + ("replay", "bfloat16", 2, "qwen3_5", 16, None), + ("replay", "bfloat16", 2, "qwen3_5", 4, None), + ("replay", "bfloat16", 0, "qwen3_5", 16, None), + ("replay", "float32", 3, "qwen4_exp", 16, None), + ("replay", "bfloat16", 3, "qwen4_exp_text", 8, None), + ("replay", "float32", 4, "qwen3_5", 4, "capacity"), + ("replay", "float32", 2, "llama", 16, "require a GDN"), + ], +) +def test_ssm_mode_validation(monkeypatch, mode, dtype, step, model, capacity, error): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) + args = make_argument_parser().parse_args( + [ + "--ssm_state_mode", + mode, + "--linear_att_ssm_data_type", + dtype, + "--mtp_step", + str(step), + "--ssm_replay_config", + json.dumps({"cache_len": capacity} if mode == "replay" else {}), + ] + ) + if error: + with pytest.raises(AssertionError, match=error): + _validate_and_normalize_ssm_config(args) + else: + _validate_and_normalize_ssm_config(args) + + +@pytest.mark.parametrize("model", ["qwen3_5", "glm5_next", "qwen4_exp", "qwen4_exp_text"]) +@pytest.mark.parametrize("mode", ["native", "replay"]) +@pytest.mark.parametrize("step", [15, 16]) +def test_gdn_kda_mtp_width_limit(monkeypatch, model, mode, step): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) + args = StartArgs(ssm_state_mode=mode, mtp_step=step) + if step <= 15: + _validate_and_normalize_ssm_config(args) + else: + with pytest.raises(AssertionError, match="GDN/KDA requires mtp_step <= 15"): + _validate_and_normalize_ssm_config(args) + + +@pytest.mark.parametrize("model", ["qwen3_5", "glm5_next", "qwen4_exp", "qwen4_exp_text"]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +@pytest.mark.parametrize("cache_len", [32, 64]) +def test_replay_large_cache_does_not_bypass_mtp_width_limit(monkeypatch, model, compute_mode, cache_len): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) + args = StartArgs( + ssm_state_mode="replay", + ssm_replay_config={"cache_len": cache_len, "compute_mode": compute_mode}, + linear_att_ssm_data_type="bfloat16", + mtp_step=16, + ) + with pytest.raises(AssertionError, match="GDN/KDA requires mtp_step <= 15"): + _validate_and_normalize_ssm_config(args) + + +def test_ssm_mode_defaults_and_removed_flag(): + parser = make_argument_parser() + assert parser.parse_args([]).ssm_state_mode == StartArgs().ssm_state_mode == "native" + assert parser.parse_args([]).ssm_replay_config == StartArgs().ssm_replay_config == {} + with pytest.raises(SystemExit): + parser.parse_args(["--enable_replayssm"]) + with pytest.raises(SystemExit): + parser.parse_args(["--ssm_state_mode", "auto"]) + + +@pytest.mark.parametrize( + "overrides,error", + [ + ({"ssm_replay_config": {"cache_len": 2}}, "power of two and at least 4"), + ({"ssm_replay_config": {"cache_len": 6}}, "power of two and at least 4"), + ({"mtp_step": -1}, "mtp_step >= 0"), + ({"linear_att_ssm_data_type": "float16"}, "FP32 or BF16"), + ], +) +def test_programmatic_replay_config_rejected_before_model_loading(monkeypatch, overrides, error): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") + args = StartArgs(ssm_state_mode="replay") + for key, value in overrides.items(): + setattr(args, key, value) + with pytest.raises(AssertionError, match=error): + _validate_and_normalize_ssm_config(args) + + +def test_replay_compute_mode_requires_replay(monkeypatch): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") + parser = make_argument_parser() + args = parser.parse_args(["--ssm_state_mode", "replay", "--ssm_replay_config", '{"compute_mode":"precompute"}']) + _validate_and_normalize_ssm_config(args) + + +@pytest.mark.parametrize( + "config,error", + [ + ([], "JSON object"), + (None, "JSON object"), + ({"projection_mode": "inline"}, "Unknown"), + ({"cache_len": True}, "capacity"), + ({"cache_len": 16.0}, "capacity"), + ({"cache_len": "16"}, "capacity"), + ({"cache_len": 128}, "capacity"), + ({"compute_mode": "auto"}, "compute_mode"), + ({"fold_method": "raw"}, "Unknown"), + ], +) +def test_invalid_replay_config(config, error): + args = StartArgs(ssm_state_mode="replay", ssm_replay_config=config) + with pytest.raises(AssertionError, match=error): + _validate_and_normalize_ssm_config(args) + + +@pytest.mark.parametrize("mode", ["native"]) +def test_replay_config_rejects_other_modes(mode): + args = StartArgs(ssm_state_mode=mode, ssm_replay_config={"cache_len": 16}) + with pytest.raises(AssertionError, match="requires ssm_state_mode=replay"): + _validate_and_normalize_ssm_config(args) + + +def test_partial_replay_config_defaults_and_serialization(monkeypatch): + from dataclasses import asdict + + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") + args = make_argument_parser().parse_args( + ["--ssm_state_mode", "replay", "--ssm_replay_config", '{"compute_mode":"precompute"}'] + ) + _validate_and_normalize_ssm_config(args) + assert args.ssm_replay_config == {"cache_len": 16, "compute_mode": "precompute"} + original = StartArgs(ssm_state_mode="replay", ssm_replay_config=args.ssm_replay_config) + restored = StartArgs(**json.loads(json.dumps(asdict(original)))) + _validate_and_normalize_ssm_config(restored) + assert restored.ssm_replay_config == args.ssm_replay_config + assert StartArgs().ssm_replay_config == {} + + +@pytest.mark.parametrize( + "option,value", + [ + ("--replayssm_cache_len", "16"), + ("--replayssm_projection_mode", "inline"), + ("--ssm_replay_config", "{invalid}"), + ], +) +def test_removed_flags_and_malformed_json(option, value): + with pytest.raises(SystemExit): + make_argument_parser().parse_args([option, value])