MAX 中 Qwen3.5 混合注意力架构（max.pipelines.architectures.qwen3_5）深入解析

MAX 中 Qwen3.5 混合注意力架构max.pipelines.architectures.qwen3_5深入解析【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo本文以 MAX 仓库中的max.pipelines.architectures.qwen3_5模块为研究对象系统讲解 MAX 如何为 Qwen3.5 这一线性注意力 全注意力混合架构提供端到端推理支持从模型图构建、Gated DeltaNet 线性注意力状态池管理到混合 NVFP4/FP8 量化方案与 KV 缓存规划再到多模态视觉编码器 M-RoPE与张量并行等进阶能力。读完本文你将掌握该模块的源码结构、核心配置参数及其底层实现原理能够直接在 MAX Pipeline 体系中加载并运行 Qwen3.5 系列模型。一、模块定位Qwen3.5 在 MAX 架构注册体系中的角色max/pipelines/architectures.qwen3_5是 MAX Pipeline 内置支持的众多模型架构之一。MAX 的架构体系遵循统一的注册模式每个架构模块都会注册一个SupportedArchitecture实例告知 Pipeline 系统如何加载、配置和执行某个模型家族见 pipelines.architectures.rst 的说明。在 arch.py 中Qwen3.5 以如下方式注册qwen3_5_arch SupportedArchitecture( nameQwen3_5ForConditionalGeneration, taskPipelineTask.TEXT_GENERATION, example_repo_ids[Qwen/Qwen3.5-27B, Qwen/Qwen3.8-27B], default_weights_formatWeightsFormat.safetensors, default_encodingQwen3_5Config.DEFAULT_ENCODING, # bfloat16 supported_encodingsQwen3_5Config.SUPPORTED_ENCODINGS, pipeline_modelQwen3_5Model, tokenizerQwen3_5Tokenizer, context_typeQwen3VLTextAndVisionContext, weight_adapters{WeightsFormat.safetensors: convert_qwen3_5_state_dict}, configQwen3_5Config, batchingQwen3_5BatchProcessor, multi_gpu_supportedTrue, tool_parserqwen3_5, reasoning_parserqwen3_5, memory_plannerQwen3_5MemoryPlanner, supports_device_graph_captureTrue, )关键信息包括注册名为Qwen3_5ForConditionalGeneration多模态条件生成架构名同时兼容纯文本的Qwen3_5ForCausalLM权重任务类型为文本生成PipelineTask.TEXT_GENERATION默认精度编码为bfloat16同时支持float32、float8_e4m3fn、float4_e2m1fnx2四种编码注册了配套的 Tokenizer、批处理器、工具解析器tool_parser、推理解析器reasoning_parser与内存规划器supports_device_graph_captureTrue要求Qwen3_5Model.release_warmup_state在每次图捕获预热后释放状态池槽位否则预热探测会耗尽池子arch.py 注释明确说明了这一点。二、核心架构全注意力与线性注意力Gated DeltaNet的混合Qwen3.5 的模型主体定义在 qwen3_5.py 中类Qwen3_5的 docstring 直接概括了架构本质This model uses a mix of full attention (with KV cache) and linear attention (Gated DeltaNet) layers. Every full_attention_interval-th layer uses full attention, and the rest use linear attention.即每full_attention_interval层使用一次全注意力携带 KV Cache其余层使用线性注意力Gated DeltaNet。这一混合设计在保留强表达力的同时用线性注意力层显著降低长上下文下的 KV 缓存压力。对应地模型图按层类型构建两种 BlockQwen3_5FullAttentionBlock全注意力 Transformer Block走 KV Cache 路径Qwen3_5LinearAttentionBlock线性注意力 Block走 Gated DeltaNet 路径。两者的前向结构高度对称——都遵循「输入 LayerNorm → 注意力混合器 → 残差 → 后注意力 LayerNorm → MLP → 残差」的经典 Transformer 块结构区别只在混合器前者用Qwen3_5Attention标准 GQA 注意力 门控输出后者用GatedDeltaNet状态空间模型风格的线性注意力。两种 Block 共享同一个_shard_mlp_and_norms辅助函数它按中间维切分 MLP 并复制两个 RMSNorm在单设备时各分片策略退化为恒等映射源码注释明确说明这一点。2.1 全注意力层门控输出与部分 RoPEQwen3_5FullAttentionBlock内部使用 layers/attention.py 中的Qwen3_5Attention。该实现与标准 GQA 注意力有三处关键差异源码 docstring 明确列出门控 Q 输出gated outputq_proj输出 2 倍宽度——hidden_size - num_heads * head_dim * 2按头交错排布为[head0 Q | head0 gate | head1 Q | head1 gate | ...]。每个头的注意力输出先经 sigmoid 门控再送入输出投影。注意如果错误地切成[all Q | all gate]两块能编译、能运行但会静默产出错误结果attention.py 注释特别警告了这一陷阱。部分 RoPEpartial RoPE只有partial_rotary_factor * head_dim默认 25%的维度施加旋转位置编码由partial_rotary_factor: float 0.25控制。Q/K 的 RMSNorm 使用(1 weight)偏移即weight_offset1.0的归一化约定而不是常见的(weight)形式。配置层面全注意力层直接继承自 Llama3 家族的参数num_attention_heads、num_key_value_heads、head_dim、rms_norm_eps等并额外增加attention_multiplier注意力缩放系数与attention_bias等选项。2.2 线性注意力层Gated DeltaNet 与状态池Qwen3_5LinearAttentionBlock使用 layers/gated_deltanet.py 中的GatedDeltaNet这是一种 SSM状态空间模型风格的线性注意力。它不维护传统 KV Cache而是维护两类可原地修改in-place的状态池卷积状态池conv pool形状[max_slots, conv_dim, K-1]其中conv_dim 2 * linear_key_head_dim * linear_num_key_heads linear_value_head_dim * linear_num_value_headsK为因果卷积核大小默认 4循环状态池recurrent pool形状[max_slots, num_v_heads, key_dim, val_dim]。池是可变BufferType图输入槽索引 SSM 内核通过指针运算直接在slot_idx[batch_item]处读写没有 gather/scatter 拷贝、也没有逐 step 的池分配。state_cache.py的 docstring 指出这一设计对齐了 vLLM 的selective_state_update方案并将原来逐 decode step 的num_layers * batch * 2 directions次 D2D 启动压缩为一次极小的slot_idxH2D 传输。值得注意的实现细节Qwen3_5LinearAttentionBlock有一个类级属性replay_capture: list[list[GatedDeltaReplayInputs]] | None None——当启用投机解码speculative graph时验证阶段过后会对接受的 token 前缀重放该层状态内核因此需要一个逐设备的重放输入收集器默认关闭时基础图不受影响。2.3 前向流程与层的编排Qwen3_5.__call__的前向流程见 qwen3_5.pyEmbeddingVocabParallelEmbedding将 token 映射为隐藏状态多模态嵌入合并若存在通过merge_multimodal_embeddings将视觉编码器输出按image_token_indices散射到 token 序列对应位置RoPE 频率准备M-RoPE 开启时按position_ids[3, total_seq_len]的时间/高度/宽度三维位置生成逐 token 频率表纯文本时退化为按cache_length token_idx索引静态表分层前向forward_sequential_layers依layer_types逐个执行 Block。inputs_for_layer为全注意力层构造 KV Cache 输入注意layer_idx是 KV Cache 内的顺序索引仅在全注意力层间计数与绝对层号不同为线性层构造slot_idx conv_pools recurrent_pools输出后处理_postprocess_logits经DistributedLogitsPostprocessMixin完成张量并行下 lm_head 的 allreduce 等。当use_subgraphs开启时层会按[全注意力层组, 线性注意力层组]分组为子图子图名形如qwen3_5_full_attention_block以便编译期更好地优化。此外return_hidden_states SELECTED_LAYERS时可通过target_layer_ids在指定层捕获隐藏状态——这是 DFlash 风格投机解码草稿模型读取目标模型中间层输出的钩子注释明确说明on_layer_output对两种层类型均会触发。三、配置参数Qwen3_5Config 全解析配置类Qwen3_5Config定义在 model_config.py 中继承自Llama3Config与ArchConfigWithVisionCache是一个kw_onlydataclass。它既承载混合注意力专有参数也负责从 Hugging Face 配置初始化、构造 KV 参数、估算批大小等职责。3.1 混合注意力参数参数默认值说明layer_types[]逐层注意力类型列表full_attention或linear_attention。若 HF 配置未提供则由_get_layer_types按full_attention_interval生成full_attention_interval4每 N 层使用一次全注意力默认每 4 层linear_key_head_dim128线性注意力层 Key 头维度linear_value_head_dim128线性注意力层 Value 头维度linear_num_key_heads16线性注意力层 Key 头数量linear_num_value_heads48线性注意力层 Value 头数量linear_conv_kernel_dim4线性注意力层因果卷积核大小partial_rotary_factor0.25全注意力层 head_dim 中施加旋转位置编码的比例attn_output_gateTrue全注意力层是否使用 sigmoid 输出门控mamba_ssm_dtypeDType.float32线性注意力层 SSM 计算的精度在 GatedDeltaNet 内部计算会被提升到 float32state_pool_dtypeNone线性注意力状态池的存储精度覆盖。None时两个池均以compute_dtypebfloat16存储设为float32可使投机解码的生成严格跟随非投机解码的状态轨迹代价是每请求状态内存约翻倍其中_get_layer_types的生成逻辑值得展开若 HF 配置携带layer_types则直接使用否则按(i 1) % full_attention_interval 0判定第 i 层为full_attention其余为linear_attention。3.2 多模态视觉参数纯文本检查点这些字段均为None多模态检查点Qwen3_5ForConditionalGeneration会携带参数说明vision_config视觉编码器配置VisionConfig复用 qwen3vl_moe 的视觉配置None表示纯文本模型image_token_id/video_token_id输入序列中图像/视频占位符使用的 token IDvision_start_token_id标记视觉内容开始的 token IDmrope_section多模态旋转位置编码M-RoPE的段长列表mrope_enabled的判定条件是mrope_section is not None and vision_config is not None。Qwen3.5 的视觉塔是NaViT 风格图像保持宽高比、每个合并 patch 产生一个 token——这与 Gemma4 固定 pooled patch 数不同因此estimate_vision_cache_entry_bytes无法从配置读取每图 token 数只能以图像处理器的缩放后像素上限 ÷ patch 面积 ÷ 空间合并数作为上界来估算视觉缓存条目大小。3.3 精度与初始化逻辑DEFAULT_ENCODING bfloat16SUPPORTED_ENCODINGS为{bfloat16, float32, float8_e4m3fn, float4_e2m1fnx2}compute_dtype是一个属性量化方案解析后以quant_scheme.compute_dtype为准在finalize之前则以检查点声明的 dtypedeclared_dtype兜底。这是为了避免误读存储 dtype打包 NVFP4 下为uint8state_dtype是所有池缓冲区声明者基础图、融合投机图、服务端状态缓存必须统一读取的单一属性——两个图在服务时共享同一份池分配任何不一致都会导致不可服务的产物对initialize_from_config会同时处理多模态顶层配置含text_config与纯文本配置即文本两种形态通过_get_text_config提取文本配置rope_theta与partial_rotary_factor的读取优先级为顶层text_config.rope_thetarope_parameters.rope_theta 基类默认值部分 RoPE 要求内核使用交错interleaved模式因此配置强制interleaved_rope_weightsTrue数据并行不受支持construct_kv_params中若data_parallel_degree 1会直接抛出ValueError。3.4 KV Cache 参数构造construct_kv_params只为全注意力层分配 KV Cachenum_full_attention_layers sum(1 for lt in layer_types if lt full_attention) ... return kv_cache_config.to_params( allow_kv_head_replication..., dtypecache_dtype, n_kv_headstext_config.num_key_value_heads, head_dimtext_config.head_dim, num_layersnum_full_attention_layers, devicesdevices, data_parallel_degreedata_parallel_degree, page_sizepage_size, )实现细节MHA 内核会以tile_size head_dim工作因此 KV 缓存page_size必须不小于 tile size。Qwen3.5 的head_dim256故当head_dim 128时强制page_size max(page_size, head_dim)。前向时每个全注意力层被映射到 KV Cache 内的顺序索引0、1、2……与绝对层号相互独立。线性注意力层的缓存则由独立的 conv/recurrent 状态池承担。3.5 批大小推断内存规划infer_optimal_batch_size为架构提供内存安全的默认max_batch_size每请求线性注意力状态字节数由_per_request_state_bytes()计算每层 conv_dim * (kernel-1) * dtype_bytes num_v_heads * key_head_dim * value_head_dim * dtype_bytes由于状态池是单块原地读写、无工作副本峰值占用即为max_batch × per_req字节内存预算按权重之后的一半分给状态池、其余给 KV 缓存的规则分配设备查询失败时保守回退到 32注释说明该值对 27B 模型在 H100/A100 80GB 上是安全的。四、混合量化方案NVFP4MLP FP8注意力Qwen3.5 的量化是该模块最值得研究的部分之一实现在 quantization.py 中。其背景是modelopt 将 Qwen3.8-27B 导出为quant_algo: MIXED_PRECISIONMLP 投影与lm_head是 group size 16 的 NVFP4而注意力与 Gated-DeltaNet 投影是逐张量 FP8。单个QuantConfig无法同时描述两者——把 FP8 载荷当作打包 FP4 读取会产生看似合理、实则垃圾的输出且不报错。因此该模块引入Qwen3_5QuantScheme用两个QuantConfig对象 各自作用的模块集合表示方案dataclass(frozenTrue) class Qwen3_5QuantScheme: mlp: QuantConfig | None # NVFP4gate/up/down_proj 与 lm_head attn: QuantConfig | None # 逐张量静态 FP8attention/GDN 投影 mlp_layers: frozenset[int] # 哪些层的 MLP 是 NVFP4 attn_layers: frozenset[int] # 哪些层的注意力/GDN 是 FP8 quantize_lm_head: bool # lm_head 是否 NVFP4 compute_dtype: DType # 未量化部分激活/归一化/嵌入/GDN 卷积/状态池的精度设计要点单一 QuantConfig 不动Linear已支持逐实例配置因此逐模块粒度不需要修改被二十多个架构共享的QuantConfig类型NVFP4E4M3 块缩放覆盖 16 元素NVFP4_GROUP_SIZE 16输入/权重缩放均为静态块粒度can_use_fused_mlp由检查点重新推导因为该导出中gate_proj与up_proj共享weight_scale_2与input_scale融合 gate/up 矩阵乘是精确的FP8逐张量静态 W8A8每个权重与激活一个 fp32 标量格式枚举为FBGEMM_FP8严苛的校验宁可失败不可出错parse_quant_scheme对检查点元数据逐一校验——只接受NVFP4/FP8/MIXED_PRECISION三种quant_algoNVFP4 group size 必须为 16MLP 只允许 NVFP4、注意力只允许 FP8lm_head只允许 NVFP4模块名必须匹配layers.N.(mlp|self_attn|linear_attn).proj模式MLP/注意力块部分量化同一块内只有部分投影被量化直接抛错因为 MAX 要求整块全量或全不量化。compute_dtype从embed_tokens.weight的 dtype 读取未量化权重的精度。storage_dtype()辅助函数决定某模块权重的磁盘 dtypeNVFP4 为uint8两个 4-bit 码打包一字节、逐张量 FP8 为float8_e4m3fn、未量化则保持 compute dtype。与此呼应model.py 中的_check_weights_match在load_state_dict(strictFalse)之外增设了严格门禁图需要的权重缺失、或检查点提供的量化 scale 张量_SCALE_SUFFIXES (.weight_scale, .weight_scale_2, .input_scale)未被消费时直接抛错——因为 scale 张量承载量化检查点必需的校准信息丢一个不是降级而是换了个模型。五、状态缓存与运行时输入编排5.1 GatedDeltaNetStateCacheGPU 常驻槽位池state_cache.py 实现GatedDeltaNetStateCache在构造时为每个线性注意力层预分配两组 GPUBufferconv_pool[l][max_slots, conv_dim, conv_kernel-1]rec_pool[l][max_slots, num_v_heads, key_dim, val_dim]其生命周期为claim登记请求并清零池行→slot_idx_for把批次槽位索引写入调用方持有的[max_batch_size]uint32 预分配缓冲→ 模型执行内核原地改写池图无状态输出→release释放槽位。模型加载时Qwen3_5Model.load_model还会对规划字节数与实际分配字节数做断言相等allocated self._accounted_state_bytes防止规划时读了编码的存储 dtype、分配时用了池 dtype这类曾导致批大小虚高、加载时 OOM 的历史 bug 复发。此外由于每次图捕获预热探测都会占用状态池槽位release_warmup_state会在每次(batch_size, cache_length)探测后释放槽位否则预热扫描会在正式服务前耗尽池子。5.2 运行时输入Qwen3_5Inputsmodel.py 中的Qwen3_5Inputs在Llama3Inputs基础上扩展了slot_idx逐设备[B]uint32 槽位索引conv_pools/recurrent_pools逐设备可变状态池device-major即先按设备、再按层组织request_ids用于状态缓存槽位管理的请求 IDdecoder_position_ids[3, total_seq_len]的 M-RoPE 三维位置仅当图以 M-RoPE 构建时存在。输入张量的完整次序在input_typesqwen3_5.py中定义为基础输入tokens、input_row_offsets、return_n_logits→ 信号缓冲区 → 扁平化 KV 类型 →slot_idx→ conv 池 → recurrent 池 → 视觉输入 → M-RoPE 位置。张量并行下 value 头被切分因此每个设备持有自己1/num_devices份额的头对应的池以及自己的slot_idx副本。5.3 不支持的能力硬性断言前向中对kv_collections的断言qwen3_5.py#L571-L580明确了边界不支持带缩放校准的量化 KV 缓存无缩放 FP8 缓存--kv-cache-format float8_e4m3fn可用kv_scales必须为None不支持eagle 投机解码draft_attention_dispatch_metadata必须为None。此外Qwen3_5.__init__对model_quantization_encoding是GPTQ或任意非None值时抛出NotImplementedErrorGGUFQ 亦未实现。数据并行同样被明确拒绝。六、多模态视觉编码器与 M-RoPE6.1 视觉图qwen3_5_vision多模态检查点会构建独立的视觉图qwen3_5_visionmodel.py#L385-L505其输入包括pixel_values、weights[4, seq, 1]、indices[4, seq]、rot_pos_ids、max_grid_size、grid_thw、cu_seqlens、max_seqlen等当前为单设备实现。视觉图输出 image embeddings在语言图前向中通过merge_multimodal_embeddings合并进 token 序列。Qwen3.5不使用 deepstack即不在多个 LM 深度注入中间视觉特征——这是 Qwen3VL-MoE 的特性Qwen3.5 只把编码器输出合并到嵌入层。pack_vision_inputs在流水线的 prep-ahead 窗口内把缓存未命中的图像打包上传到设备使 H2D 拷贝与上一批次重叠vision_execute运行编码器并返回嵌入纯文本/缓存命中批次则返回零行空嵌入缓冲且该缓冲被缓存复用避免每次分配。6.2 M-RoPE三维位置编码M-RoPE 开启时多模态 mrope_section非空RoPE 由Qwen3_5TextRotaryEmbedding提供位置表按 token 组织而非按位置组织freqs_cis_position_ids(position_ids)为[3, total_seq_len]的每列生成频率内核用显式位置 id 索引。文本-only 检查点没有图像拼接三个轴全部退化为默认的cache_length token_idx因此保持静态表。KV 缓存 dtype 不参与此逻辑——RoPE 施加在 K 上、值入缓存之前。七、如何在 MAX 中运行 Qwen3.5依据SupportedArchitecture注册信息与本文介绍的配置机制可以在 MAX Pipeline 中按如下方式运行仓库为只读以下均为查看与运行方式# 通过架构名与示例模型仓库运行文本生成 python -m max.pipelines.pipeline \ --model-path Qwen/Qwen3.5-27B \ --task text-generation # 多模态检查点Qwen3_5ForConditionalGeneration 形态 python -m max.pipelines.pipeline \ --model-path Qwen/Qwen3.8-27B \ --task text-generation要点提示默认编码bfloat16可显式指定float8_e4m3fnFP8 量化或float4_e2m1fnx2NVFP4 量化量化方案由检查点顶层的quantized_layers元数据自动解析KV 缓存格式支持无缩放的float8_e4m3fn线性注意力层的状态池可通过state_pool_dtype配置为bfloat16或float32模型不支持数据并行但支持多 GPU 张量并行multi_gpu_supportedTrue池与 value 头随张量并行切分到各设备权重格式为 safetensors加载时经过convert_qwen3_5_state_dict适配mtp.*投机解码头会被丢弃。八、模块源码地图以下是max.pipelines.architectures.qwen3_5模块的完整文件结构便于按需深入max/python/max/pipelines/architectures/qwen3_5/ ├── arch.py # SupportedArchitecture 注册Qwen3_5ForConditionalGeneration ├── model.py # Qwen3_5Model语言图 视觉图构建、权重校验、状态缓存装配 ├── qwen3_5.py # Qwen3_5 模型主体混合 Block 定义与前向流程 ├── model_config.py # Qwen3_5Config混合注意力/多模态/量化配置与初始化 ├── quantization.py # Qwen3_5QuantSchemeNVFP4(MLP) FP8(注意力) 混合方案解析 ├── state_cache.py # GatedDeltaNetStateCacheGPU 常驻状态槽位池 ├── batch_processor.py # 批处理器含 slot_idx/池的准备 ├── memory_planner.py # 内存规划器 ├── tokenizer.py # 分词器 ├── reasoning.py # 推理解析器注册 qwen3_5 ├── tool_parser.py # 工具调用解析器注册 qwen3_5 ├── vision_packing.py # 图像打包pack_uncached_images ├── weight_adapters.py # 状态字典适配convert_qwen3_5_state_dict ├── mtp.py # MTP 投机解码草稿头 └── layers/ ├── attention.py # Qwen3_5Attention门控输出 部分 RoPE ├── gated_deltanet.py # GatedDeltaNet 线性注意力 ├── text_rotary.py # Qwen3_5TextRotaryEmbeddingM-RoPE └── visual_transformer.py# VisionTransformer 视觉编码器总结max.pipelines.architectures.qwen3_5展示了 MAX 对新一代混合注意力 LLM 的完整工程化支持全注意力层与 Gated DeltaNet 线性注意力层在同一个图中协同运行分别由 KV Cache 与原地更新的状态池承载缓存职责混合量化以双QuantConfig方案精确描述 NVFP4/FP8 并存检查点并以严格校验杜绝静默错误多模态能力通过 NaViT 风格视觉编码器与 M-RoPE 三维位置编码集成。理解该模块即可举一反三地掌握 MAX 架构注册、图构建、内存规划与量化解析的通用机制。【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考