MAX Pipelines 的 LFM2 架构支持混合全注意力 短卷积解码器实现解析【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo导读本文围绕 MAX 推理框架Modular Platform 的一部分即当前仓库max/目录所对应的 Python 服务端中max.pipelines.architectures.lfm2模块展开剖析它如何将 LiquidAI 的 LFM2 系列模型如LiquidAI/LFM2.5-350M接入 MAX Pipelines 推理管线。读者将了解 LFM2 的「全注意力 短卷积」混合架构在 MAX 中的图构建、权重适配、分页 KV Cache 与卷积状态缓存conv-state cache实现以及如何在仓库中定位与验证这套支持代码。文档定位一个由 automodule 驱动的 API 模块页关联文档 max/python/docs/pipelines.architectures.lfm2.rst 是 MAX Python 文档体系中max.pipelines.architectures命名空间下的一个 Sphinx 模块页。其正文只有一条automodule指令.. automodule:: max.pipelines.architectures.lfm2 :members: :imported-members: :show-inheritance:也就是说该文档的实质性内容完全来自被自动索引的 Python 模块max.pipelines.architectures.lfm2。docs/目录下还包含pipelines.architectures.rst、pipelines.rst等一系列同构页面如pipelines.architectures.llama3.rst、pipelines.architectures.qwen3.rst共同构成架构层 API 文档。因此理解本文主题的正确路径是直接阅读模块源码目录 max/python/max/pipelines/architectures/lfm2/该目录包含arch.py— 架构注册SupportedArchitecturemodel.py— 管线模型LFM2Model与卷积状态缓存ConvStateCachelfm2.py— 图级模型LFM2、解码层LFM2DecoderLayer、短卷积LFM2ShortConv、MLPfull_attention.py— 全注意力块LFM2FullAttentionmodel_config.py— 配置类LFM2Configbatch_processor.py— 批处理器LFM2BatchProcessorweight_adapters.py— safetensors 权重名映射__init__.py— 公共导出ConvStateCache、LFM2Config、LFM2Inputs、LFM2Model、lfm2_arch架构注册LFM2 如何被 MAX Pipelines 识别在 arch.py 中模块通过SupportedArchitecture将 LFM2 注册为一种因果语言模型架构字段值说明nameLfm2ForCausalLM与 Hugging Face 模型 config 中的architectures字段对应taskPipelineTask.TEXT_GENERATION文本生成任务default_encoding/supported_encodingsfloat32/{float32, bfloat16}默认 FP32 推理支持 BF16example_repo_ids[LiquidAI/LFM2.5-350M, LiquidAI/LFM2.5-350M-Base]官方示例仓库 IDtokenizerTextTokenizer文本 tokenizercontext_typeTextContext文本上下文default_weights_formatWeightsFormat.safetensors默认 safetensors 权重required_argumentsallow_safetensors_weights_fp32_bf16_bidirectional_castTrue、trust_remote_codeTrue加载 LFM2 权重时必须开启的两个参数multi_gpu_supportedFalse当前不支持多 GPU 数据并行weight_adapterssafetensors →convert_lfm2_safetensor_state_dict权重名转换configLFM2Config配置类batchingLFM2BatchProcessor批处理memory_plannerPagedMemoryPlanner分页内存规划supports_overlap_scheduler/supports_device_graph_captureFalse不支持重叠调度与设备图捕获从源码结构看MAX Pipelines 通过arch_lookup见 max/python/docs/pipelines.lib.arch_lookup.rst 对应的max.pipelines.lib模块与各架构arch.py中的注册信息根据 Hugging Faceconfig.json的architectures字段自动选择Lfm2ForCausalLM实现。__init__.py通过from .arch import lfm2_arch导出lfm2_arch并公开LFM2Model、LFM2Config、LFM2Inputs、ConvStateCache这正是 Sphinx 文档页:members:会索引的核心 API。混合架构核心全注意力层 短卷积层LFM2 的关键设计是按层混合layer_types列表逐层标记每层是full_attention还是卷积层。这一结构在 lfm2.py 的LFM2与LFM2DecoderLayer中体现if self.layer_type full_attention: self.self_attn LFM2FullAttention(...) self.conv None else: self.conv LFM2ShortConv(config, linear_cls) self.self_attn None全注意力层使用LFM2FullAttention继承自 Qwen3.5 风格 GQA 实现被 vendor 到 full_attention.py非注意力层使用LFM2ShortConv短卷积对局部窗口内的 token 做因果卷积convolution over the sequence dimension从而在保持全局建模能力的同时降低注意力层数量带来的开销。LFM2在构造时统计layer_types中非full_attention的层数num_conv_layersKV 索引kv_idx只对全注意力层递增卷积层不占用 KV cache 页。短卷积实现LFM2ShortConvLFM2ShortConv 的图级实现要点in_proj将隐藏状态投影为 3 份b、c、xv计算bx b * xv类似门控线性单元conv_weight形状为[hidden, 1, kernel_size]kernel_size来自配置conv_L_cache默认 3通过ops.while_loop逐 token 滑动窗口每个 token 从对应请求的[1, hidden, K]状态中取出、滑入新 tokenstate_r_next ops.concat((state_r[:, :, 1:], bx_t_k), axis2)再用scatter_nd写回状态栈卷积输出conv_out sum(state_r_next * conv_w, axis2)最终y c * conv_out经过out_proj输出并返回更新后的new_state。代码注释特别指出scatter_nd是 GPU 友好的ops.scatter会经由 CPU round-trip见 TODO(GEX-2197)因此在循环内更新状态栈使用scatter_nd。这是从源码注释可确认的工程取舍。全注意力块LFM2FullAttentionfull_attention.py 实现了 Qwen3.5 风格的全注意力路径与qwen3_5_moe.layers.attention逻辑一致模块头部 docstring 明确说明该实现被 vendor 于此避免 LFM2 依赖整个 qwen3_5_moe 架构包per-head Q/K RMSNormq_norm/k_normQwen3 风格可选输出门attn_output_gate开启时q_proj输出维翻倍checkpoint 按「逐头交错」布局存放[Q_h0, Gate_h0, Q_h1, Gate_h1, ...]因此必须 reshape 为[seq_len, n_heads, 2*head_dim]后沿最后一维切分而非扁平切分门在注意力输出后以attn_output * sigmoid(gate)形式生效非 silu非融合 KV 路径门开启时fused_qkv_ragged_matmul无法处理翻倍的 Q 维故改用matmul_kv_cache_raggedconcat(k_proj.weight, v_proj.weight)写入分页 cache并配合rms_norm_key_cache对缓存中的 K 做 per-head normpartial RoPEfreqs_cis.shape[-1] head_dim * partial_rotary_factorfused_qk_ragged_rope只旋转前导维度因果 Flash Attentionflash_attention_ragged支持local_window_size默认 512与MHAMaskVariant.CAUSAL_MASK张量并行实现Shardable支持 replicate 与 tensor_parallel 两种分片策略门开启时q_proj使用gate_up分片保证每个分片拿到[q_shard | gate_shard]o_proj使用head_aware_columnwiseKV 投影使用rowwise。需要注意该支持是图内张量并行与arch.py中multi_gpu_supportedFalse并不矛盾后者针对数据并行多副本。配置类 LFM2Config从 Hugging Face config 到图构建参数model_config.py 中LFM2Config(Llama3Config)新增字段字段默认值来源HF config 键说明layer_types[]layer_types每层类型full_attention或卷积层conv_L_cache3conv_L_cache卷积缓存窗口长度kernel sizeconv_biasFalseconv_bias卷积是否带偏置norm_eps1e-5norm_eps归一化 epsiloninitialize_from_config从 HFAutoConfig读取这些字段并处理rope_parameters.rope_theta默认DEFAULT_ROPE_THETA 10000.0与可选 rope 字段缺失的兼容_ensure_optional_rope_fields。值得注意的两个兼容细节均有源码注释佐证norm_epsvsrms_norm_epsLFM2 将归一化 epsilon 存在norm_eps而非 LLaMA 风格的rms_norm_eps。finalize()先以norm_methodlayer_norm调用父类以跳过对rms_norm_eps的属性读取随后再改回rms_norm_eps并显式赋值self.rms_norm_eps norm_epstie_embeddingvstie_word_embeddingstransformers ≥ 5 将 LiquidAI 的自定义tie_embedding键并入标准tie_word_embeddings代码先读标准键再回退到旧键。_resolve_intermediate_size还处理block_auto_adjust_ff_dimSwiGLU FFN 有两个门控投影有效宽度按 2/3 折算与 LLaMA 系列同约定并依次应用block_ffn_dim_multiplier与block_multiple_of对齐。LFM2图模型lfm2.py的组件包括embed_tokens嵌入、rope由create_rope_embedding创建interleaved_rope_weightsFalse、layers层列表、normRMSNorm与lm_head当tie_word_embeddings为真时lm_head共享embed_tokens.weight。前向流程为嵌入 → 逐层残差 注意力/卷积 SwiGLU FFN→logits_postprocess含可选的 logits 缩放→ 拼接各卷积层新状态。推理状态管理ConvStateCache 与卷积状态生命周期卷积层没有 KV cache但需要跨 token 维护滑动窗口状态因此 MAX 在管线侧实现了专门的 ConvStateCache每个 slot 为每层分配Buffer.zeros([1, hidden_size, conv_L_cache], dtype, device)claim(request_id)为请求申请 slot无空闲 slot 时抛RuntimeError(No free LFM2 conv-state slots.)release(request_id)归还 slotget_states(request_ids)单请求时零拷贝直接返回 slot buffer多请求N 1时按层沿 batch 维拼接为[N, hidden, kernel]update_states(request_ids, new_states)执行图输出新状态后按请求写回各自 slot。拼接/切分通过 numpy round-trip 完成_cat_buffers/_split_buffer_dim0其中对 bfloat16 的处理非常细致numpy 不支持 bfloat16DLPack 转换会抛RuntimeError: Unsupported dtype in DLTensor因此先将 bf16 视图为同字节宽的uint16走 numpy再视图回原 dtype。由于卷积状态本身很小每 slot 为hidden * kernel这种 round-trip 开销可接受——这是源码 docstring 明确给出的设计权衡。LFM2Modelmodel.py继承Llama3ModelBase_create_model_config调LFM2Config.initialize后finalize_build_graph_for_compile构建名为lfm2的Graph输入为 tokens、input_row_offsets、return_n_logits、展开的 KV 输入、以及每个卷积层一个[conv_batch, hidden, conv_L_cache]状态输入见LFM2.input_types输出包含 logits数量由_num_logit_outputs决定依据ReturnLogits/ReturnHiddenStates与各卷积层新状态execute将图输出切成 logits 与新状态两部分新状态回写ConvStateCacherelease(request_id)释放卷积状态 slot。批处理与权重适配LFM2BatchProcessor 继承Llama3BatchProcessor在基类 ragged batching 基础上叠加卷积状态 slotbind_conv_cache由LFM2Model.__init__在构造后调用将ConvStateCache注入批处理器prepare_initial_token_inputs先调父类准备 KV/输入再对扁平化后的请求列表逐个claim最后get_states组装LFM2InputsLFM2Inputs在Llama3Inputs基础上增加conv_states: list[Buffer]与request_ids: list[RequestID]buffers属性把卷积状态追加到父类 buffer 列表之后保证与图输入顺序一致。weight_adapters.py 定义了 HF checkpoint → MAX 命名映射LFM2_SAFETENSOR_MAPPING ( (model.embed_tokens, embed_tokens), (model.embedding_norm, norm), (model.layers, layers), (self_attn.out_proj, self_attn.o_proj), (self_attn.q_layernorm, self_attn.q_norm), (self_attn.k_layernorm, self_attn.k_norm), (conv.conv.weight, conv.conv_weight), (conv.conv.bias, conv.conv_bias), )convert_lfm2_safetensor_state_dict对 state dict 中每个键按序做字符串替换如self_attn.q_layernorm→self_attn.q_norm、conv.conv.weight→conv.conv_weight完成权重名对齐。该适配器与arch.py中weight_adapters{WeightsFormat.safetensors: ...}对应是_build_graph_for_compile中model.load_state_dict(state_dict, ...)之前必经的一步。如何在仓库中验证与运行单元/集成证据LFM2 架构测试位于 max/tests/ 下模型集成测试可参考max/python/test/与max/tests/integration/中的 pipelines 测试组织方式当前仓库以 Bazel 管理相关测试目标定义于对应 BUILD.bazel。注册证据max/python/max/pipelines/architectures/__init__.py与all_arches.bzl中均列出lfm2确认其参与架构自动发现。实际使用前提加载 LFM2 权重必须满足arch.py中required_arguments的两项——trust_remote_codeTrueLiquidAI 模型包含自定义建模代码与 safetensors 的 FP32/BF16 双向转换开关且当前注册信息显示不支持多 GPU 数据并行multi_gpu_supportedFalse、不支持重叠调度与设备图捕获。文档入口完整的架构级 API 索引见 max/python/docs/pipelines.architectures.rst模块页即本文开头的pipelines.architectures.lfm2.rst服务部署与运行参数请参考 max/docs/serve/ 与 max/docs/get-started.mdx。小结max.pipelines.architectures.lfm2为 LFM2 系列混合架构模型提供了一套完整的 MAX Pipelines 支持通过SupportedArchitecture注册让Lfm2ForCausalLM可被自动发现LFM2/LFM2DecoderLayer按layer_types混合实例化 Qwen3.5 风格全注意力块与短卷积块ConvStateCache与LFM2BatchProcessor负责卷积状态在分页推理与 ragged batching 下的生命周期管理LFM2Config与convert_lfm2_safetensor_state_dict分别完成 HF config 与权重名的兼容对齐。这套实现既复用了 LLaMA 3 管线基座又以 vendor 方式引入 Qwen3.5 注意力逻辑是研究「如何在大型推理框架中接入混合注意力新架构」的典型样本。【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考