vLLM-Omni 扩散模型性能模式实战从测量循环到算子融合、USP 与解耦的完整优化指南【免费下载链接】vllm-omniA framework for efficient model inference with omni-modality models项目地址: https://gitcode.com/GitHub_Trending/vl/vllm-omni本指南以 vLLM-Omni 仓库中production-add-diffusion-model技能的performance-patterns.md为骨架系统讲解把一个 Day-0 可跑的扩散模型DiT推向可测量、可验收、可上线的过程中必须掌握的性能方法论与代码模式先讲如何建立正确的测量循环再逐一拆解更快的 BF16 注意力、稀疏注意力、融合优先的算子模式Q/K RMSNormRoPE、融合 QKV、Gate-upSwiGLU、AdaLN、冗余计算消除、USP 序列并行通信优化、文本编码器与 VAE 解耦最后给出严格的性能验收清单。读完本文你将掌握一套先固定工作负载、一次只做一个优化、以端到端指标为准的工程化优化流程并能在仓库源码中定位每一类优化对应的共享算子与实现证据。测量循环一切优化都从固定工作负载开始性能文档的第一条铁律是不要以优化清单开头要以固定的正确性与性能工作负载开头。优化前必须先冻结一个可复现的 workload 描述它至少应包含model/checkpoint revision prompt/media hashes seed dimensions/frames scheduler/sigma/steps/guidance task dtype backend topology hardware只有模型/权重版本、输入哈希、随机种子、尺寸/帧数、调度器/sigma/步数/引导强度、任务类型、dtype、注意力后端、并行拓扑、硬件全部固定才能保证前后对比有意义。需要捕获的指标族阶段耗时stage timeload、encode、conditioning、denoise、VAE decode/postprocess 各阶段时间GPU tracekernel 时长与个数、launch 间隔、显存分配、H2D 拷贝、collective集合通信开销CPU trace预处理、逐层同步、Python dispatch、临时文件内存每 rank 的 HBM 峰值、进程树的 host PSS输出链路device 准备、D2H/IPC、payload 字节数、codec 墙钟与进程 CPU、事件循环等待、客户端物化、以及 signed residual不能被掩盖到其他阶段的剩余耗时服务侧队列时间、吞吐、错误、请求混合、分布采样。A/B 方法论对于固定工作的单请求 A/B先跑一次显式 warmup 策略再做至少三次测量重复保留每一次原始 run报告 median/mean 并附范围。不要从三次样本推导 p95/p99尾延迟必须声明到达模型arrival model并采集足够多的成功请求才能支撑所声称的分位点与置信度。一次只实现一个优化重跑 parity同时报告局部 kernel 与端到端两方面的 delta——一个 kernel 赢了但端到端指标毫无变化不能作为生产速度声明。文档还强调优先使用目标版本已暴露的 pipeline profilerPyTorch/CUDA/ROCm/NPU/XPU 各自使用平台配套的 profiling 工具不要把每个平台都强行塞进 CUDA-only 的工作流。更快的 BF16 注意力共享角色感知 Attention 层与后端资质共享注意力接线Shared attention wiring不要每个模型自研一套注意力实现而是使用vllm_omni.diffusion.attention.layer中共享的Attention层并传入真实的局部TP 切分后head 数与 QKV 布局。标准接线如下import torch from vllm.model_executor.layers.linear import ( QKVParallelLinear, RowParallelLinear, ) from vllm_omni.diffusion.attention.layer import Attention class SelfAttention(torch.nn.Module): def __init__( self, hidden_size, head_dim, num_heads, num_kv_heads, quant_configNone, prefix, ): super().__init__() self.head_dim head_dim self.qkv_proj QKVParallelLinear( hidden_sizehidden_size, head_sizehead_dim, total_num_headsnum_heads, total_num_kv_headsnum_kv_heads, biasFalse, quant_configquant_config, prefixf{prefix}.qkv_proj, ) self.attention Attention( num_headsself.qkv_proj.num_heads, num_kv_headsself.qkv_proj.num_kv_heads, head_sizehead_dim, softmax_scalehead_dim**-0.5, causalFalse, qkv_layoutBSND, roleself, prefixprefix, ) self.out_proj RowParallelLinear( num_heads * head_dim, hidden_size, biasFalse, input_is_parallelTrue, quant_configquant_config, prefixf{prefix}.out_proj, ) def forward(self, x, attn_metadataNone): qkv, _ self.qkv_proj(x) q_size self.qkv_proj.num_heads * self.head_dim kv_size self.qkv_proj.num_kv_heads * self.head_dim q, k, v qkv.split([q_size, kv_size, kv_size], dim-1) q q.unflatten(-1, (self.qkv_proj.num_heads, self.head_dim)) k k.unflatten(-1, (self.qkv_proj.num_kv_heads, self.head_dim)) v v.unflatten(-1, (self.qkv_proj.num_kv_heads, self.head_dim)) out self.attention(q, k, v, attn_metadata) out out.flatten(-2) out, _ self.out_proj(out) return out关键点在调用Attention之前按官方顺序先加 Q/K norm 与 RoPE见下一节融合模式保持融合 QKV 与输出投影的 checkpoint 映射完整Attention层内部会依据role与全局 diffusion 配置解析后端。从 layer.py 的实现可以看到它支持role、role_category、qkv_layout、scatter_idx/gather_idxUlysses 注意力、skip_sequence_parallel跳过序列并行通信、custom_attention模型自有 kernel、allow_fp32_fallbackFP32 输入自动回退 SDPA等参数并且当ring_degree 1时会自动挂载RingParallelAttention。后端资质Backend qualification共享Attention类并不自动保证每个后端对 packed 样本都安全或更快。对每个角色self、cross-attention 类别、token refiner 等都要独立记录实际选中的后端从日志/trace 中确认支持的 dtype、head size、QKV 布局、mask/metadata、packed varlen 行为平台与计算能力TP/SP/Ring/AllGather-KV 限制fallback 行为以及 fallback 是否会改变声称的性能。验证时要使用长度不等的 packed 样本与 dense SDPA 对比核对cu_seqlens、最大序列长度、padding 排除、模态边界、CFG 归属与输出切片。不要传了一个attn_mask就假设它被消费了——要检查后端契约与 trace。从 selector.py 可以看到后端的解析优先级是attention_config.per_role[role]→per_role[role_category]→attention_config.default→ 平台默认。TRTLLM、Flash Attention、cuDNN、NPU varlen 等后端是相互独立的资质行Ring 与 hybrid 路径需要各自的边界 parityAllGather-KV 与特定后端组合可能直接 fail fast例如 layer.py 中TRTLLM_ATTN与allgather_degree 1组合会抛错这些行为要写进文档保持可见。稀疏注意力是后端/模型/平台优化不是低精度稠密注意力的同义词稀疏注意力只有在稠密 BF16 parity 达成之后才能开启。对每个候选任务/形状/拓扑必须声明稀疏策略、选定的后端/kernel、block/head 维度、支持的硬件、外部 kernel/包版本与 fallback证明稀疏 kernel 真的执行了并报告实际稀疏率——一个 enabled 标志不够用同种子对比逐 block 输出、denoise 轨迹与最终产物与稠密注意力增加模态专属指标视频的时序一致性与音画同步、音频的频谱/音质、以及相关的 prompt/图像指标报告注意力时间与完整端到端时间包含预处理/index-build 开销、峰值内存与精确的稠密对照组逐任务、逐形状验证不要把一个 T2V 结果外推到 I2V/Ref2V也不要把 NPU 的稀疏支持外推到 CUDA/ROCm/XPU。如果 Ring 绕过了选定后端或某个稀疏设置会被静默忽略就应拒绝该组合。只有当 recipe 把该行标注为limited且不声称稀疏加速时才允许保留可见的稠密 fallback。融合优先的算子模式优先共享算子 显式守卫 原生 fallback融合的准则是优先使用当前仓库的共享/自定义算子在算子、block、轨迹、产物四个层面证明 parity再用 trace 证明 launch 次数/字节数更少或延迟更低。Q/K RMSNorm RoPE先给出通用跨平台基线两个独立分派算子不是自动一个融合 kernelfrom vllm_omni.diffusion.layers.norm import RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk self.norm_q RMSNorm(head_dim, epseps) self.norm_k RMSNorm(head_dim, epseps) # Example only: set these booleans from the official checkpoint contract. use_neox_layout False cos_sin_are_half_dim True self.rope RotaryEmbedding( is_neox_styleuse_neox_layout, half_head_dimcos_sin_are_half_dim, ) q self.norm_q(q) k self.norm_k(k) q, k apply_rope_to_qk(self.rope, q, k, (cos, sin))注意RotaryEmbedding(is_neox_style...)的语义False表示 interleaved/GPT-J 风格偶数奇数维配对True表示 NeoX 风格前半/后半配对不要假设默认是 NeoX。真实代码中要把示例布尔值替换为官方 checkpoint 契约。部分旋转维partial rotary通过旋转前缀 直通后缀拼接实现如果 3D cos/sin 含 batch 轴当前共享路径假设值在 batch 上共享逐样本不同位置需要经过验证的表示/路径。对 packed 的[tokens, heads, head_dim]张量配合非 interleaved 的[cos | sin]表使用共享融合边界from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope q, k fused_qk_norm_rope( q, k, self.norm_q.weight, self.norm_k.weight, rope_table, self.norm_q.variance_epsilon, )从 fused_qk_norm_rope.py 的实现可以看到CUDA 快路径当前特化 BF16head_dim128、rotary_dim96MiniMax-H3 几何契约interleaved 模式支持任意偶数rotary_dim head_dim 256其他合法输入走 eager 实现_apply_rope_table支持 interleaved 与 half-split 两种配对语义分别对应 Boogu-Image 与 MiniMax-H3AscendNPU路径通过mindiesd.rotary_position_embedding/torch_npu.npu_rotary_mul组合实现同样限定 MiniMax-H3 几何函数以自定义算子注册torch.ops.vllm_omni.fused_qk_norm_rope并且提供VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS环境变量token 数低于交叉点时 fused 路径是 host-bound 的反而更慢Boogu-Image 在单张 H200 上实测默认 2048 tokens0表示总是融合。融合 QKV 投影使用上文的QKVParallelLinear。资质要点num_heads与num_kv_heads必须能被 TP 整除使用局部TP 切分后split 尺寸TP 之后绝不复用全局计数在交给 vLLM loader 前完成 Q/K/V checkpoint 顺序转换所有 Q、K、V 分片都必须存在严格加载缺失即失败启用 FP8 等方法时核对量化 scale/loader 映射。Gate-up 投影 融合 SwiGLUfrom vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, RowParallelLinear, ) self.gate_up_proj MergedColumnParallelLinear( input_sizehidden_size, output_sizes[intermediate_size, intermediate_size], biasFalse, quant_configquant_config, prefixf{prefix}.gate_up_proj, ) self.act_fn SiluAndMul() self.down_proj RowParallelLinear( input_sizeintermediate_size, output_sizehidden_size, biasFalse, input_is_parallelTrue, quant_configquant_config, prefixf{prefix}.down_proj, ) gate_up, _ self.gate_up_proj(x) x self.act_fn(gate_up) x, _ self.down_proj(x)checkpoint 的gate_proj与up_proj两个分片都要映射进 merged 参数且两者都要求存在。文档特别提醒merged 投影后跟F.silu(gate) * up与SiluAndMul融合激活 kernel不是同一个声明——要通过 trace 确认实际用的是哪一个。AdaLN只有当共享层的数学/返回契约与官方模型一致时才使用共享层from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm self.adaln AdaLayerNorm( hidden_size, elementwise_affineFalse, eps1e-6, ) x self.adaln(x, scale[:, None, :], shift[:, None, :])AdaLayerNormZero、AdaLayerNormZeroSingle及连续变体包含不同的投影、广播与 tuple 输出不能互相 drop-in 替换。从 adalayernorm.py 可以看到AdaLayerNorm在 NPU 上若有mindiesd会走layernorm_scale_shift融合原语否则 fallback 到torch_npu.npu_layer_norm_eval而 CUDA/HIP/XPU 走forward_nativelayernorm(x) * (1 scale) shift——因此它是一个带 NPU 快路径的共享层但不要称它为通用的融合 AdaLN kernel。对时间步集合有限的 sigma schedule可以先 profile 一下每个请求预计算一次 timestep/AdaLN 投影是否值得缓存 key 必须包含精确的 schedule/config/device/dtype 且保持有界。布局、残差与其他候选profiling 之后继续检查bias/dropout/add 或 gated residual 链normlinear 或 normmodulation 边界注意力周围的重复 transpose/reshape/contiguousVAE 的 normalization/activation/convolution 边界真正使用 MoE 的模型的 MoE routing 与 fused experts。融合 residual-plus-norm 或 modulation 链时保留未融合的物化materialization边界如果 residual 在 RMSNorm 前以 BF16 存储就要在融合 kernel 内先舍入到 BF16 再提升回计算精度直接使用未舍入的 FP32 临时量即使看起来对也会改变 denoise 轨迹。DLO 的物化/释放 hooks、进程组集合通信、请求缓存迁移与清理等生命周期操作要放在编译/融合区域之外除非生命周期安全性被显式测试覆盖。只编译稳定区域并保留 eager fallback。消除冗余计算提升请求不变量Hoist request invariants审计 denoise 循环中只依赖请求、不依赖当前时间步的工作把它们提升到循环之前# Before the denoise loop. condition self.condition_proj(text_embeddings) refined_condition self.token_refiner(condition, packed_metadata) rope_tables self.build_rope_tables(position_ids) attn_metadata self.build_attention_metadata(layout) for step, sigma in enumerate(sigmas): timestep_state self.prepare_timestep_state(sigma) noise self.transformer( latents, refined_condition, rope_tablesrope_tables, attn_metadataattn_metadata, timestep_statetimestep_state, ) latents self.scheduler_step(noise, sigma, latents)候选包括prompt conditioning、token-refiner 输出、reference 编码/扫描、position IDs/RoPE 表、packed masks/metadata、静态 CFG 分支、有限 schedule 的 AdaLN 投影。先证明它们对所有官方任务、LoRA/adapters、prompt-update 特性、形状与 schedule 都不变在连续 batching 中请求专属结果要存在请求状态上而不是 pipeline 单例上。消除逐层同步避免在每层内部做 device-to-host 标量提取如.item()# Bad: forces a GPU/CPU synchronization in every attention layer. max_len int(cu_seqlens[-1].item()) # Better: compute validated host metadata once during request packing and pass it. packed PackedLayout( cu_seqlenscu_seqlens, max_seqlenmax(sample_lengths), ) for block in self.blocks: hidden block(hidden, packed_layoutpacked)用 trace 验证同步确实消失。不要为了去掉.item()而把正确的动态值换成过期的常量。移除未使用的 mask 与分配如果 packed 注意力消费cu_seqlens而忽略 dense mask就不要构造 dense mask。安全时复用不可变请求元数据与预分配缓冲。优先用直接索引的 scatter/gather 替代构造大型临时张量——前提是重复索引语义与梯度推理时通常禁用保持正确。任何跨请求缓存都需要有界生命周期且 key 必须完整覆盖 shape、layout、device、dtype、schedule、task、adapter/LoRA 以及任何会改变张量的值否则保持请求级作用域。VAE 与 reference 工作profile 重复的 VAE decode/postprocess、重复的 reference 解码/扫描、不必要的 host-device 往返与整帧物化。启用 VAE tiling、slicing、patch 并行或按需暂存时必须配套 seam、顺序、颜色/范围与内存 parity 测试。USP 通信优化从能跑到通信被优化USP 能跑不等于 USP 通信被优化了。对 Ulysses、Ring、hybrid 与 AllGather-KV 要分开profilecollective 的数量、字节数、时长、所在 stream 以及与计算的 overlapQ/K/V 与输出的 all-to-all 边界pack/unpack、padding、metadata 广播、transpose 与 contiguous 拷贝head 整除性或高级 uneven-head 支持同一 workload/互连上从 1 到 N 设备的扩展效率长度不等 packed 样本的负载均衡。_sp_plan的 split/gather 边界应声明在有意义的模块输出处而不是任意的 Python 语句packed metadata 要按与 token 相同的 ownership 模型做 shard/gather。加速 Ulysses 传输--ulysses-a2a-permute当目标版本提供加速的 Ulysses 传输如--ulysses-a2a-permute时要把它与普通 Ulysses分开资质。从 ulysses.py 的实现可见加速路径的启用条件是严格 scatter-heads/gather-sequence 布局scatter_idx 2 and gather_idx 1配置项ulysses_a2a_permute定义于 omni_config.py。实践要点证明严格布局确实选中了快路径记录一次性 JIT/readiness 开销在 CUDA graph 捕获前对最大 workspace 形状做 prewarm让 grow-only workspace 保持在单条 stream 上在进程组销毁前执行 cleanup。对 GQA验证Q_heads % KV_heads 0先 pad KV heads再由原始比例推导 Q 的 padding。该文件还内置了一个警告机制当 advanced_uaa 的 GQA padding 把 Q head 数膨胀超过 1.5 倍时最坏情况 MQA 从 H 膨胀到 N*H会一次性告警提示用户选择更合适的ulysses_degree。同时要同时测试 contiguous 与 row-strided 两种 QKV 布局快拷贝路径必须保持逐位一致与 packed 边界。复制型 cross-attention 的 skip_sequence_parallelcross-attention 有时使用复制的文本 K/V只有视频/音频 Q 做序列 shard。这种特定场景下skip_sequence_parallelTrue可以避免错误或浪费的 Q/K/V 重分布self.cross_attention Attention( num_headslocal_heads, num_kv_headslocal_kv_heads, head_sizehead_dim, softmax_scalehead_dim**-0.5, causalFalse, rolecross, skip_sequence_parallelTrue, prefixf{prefix}.cross_attn, )只有在证明 K/V 确为复制、Q/输出 ownership 正确之后才使用——否则它会禁用必要的通信。报告该优化是否减少了 collective 字节/时间并改善端到端扩展。文本编码器与 VAE 解耦先测量再架构解耦disaggregation是一个需要实测的架构决策不是默认方案。先收集五问五测问题需要测量的量阶段是否足够大目标并发与形状下的 encode/decode 占比传输是否划算tensor/media 字节、序列化、connector 延迟与带宽是否有复用prompt/reference 复用率、可缓存性、fan-out能否独立扩展各阶段的队列饱和度与资源利用率契约是否稳定shape、dtype、mask/tag/position 元数据、顺序文本编码器候选定义显式 tensor 契约hidden states、masks、tags 或 position IDs、dtype、shape、task/checkpoint revision、request ID。支持禁用本地编码器避免 DiT 意外重算它。测试所有任务、prompt 长度、部分 TP 组成员、重试、背压、abort、connector 清理与独立副本扩展。VAE 候选如果 encode 与 decode 契约都存在要分开定义latent shape、scale/shift、dtype/range、tiling/chunk 顺序、输出媒体顺序与阶段元数据。用 latent/media 传输量对比本地 decode 成本。测试 tiling 接缝、相关场景的 RNG/确定性、chunk 顺序、codec/postprocess、背压、abort 与 worker 失败。文档明确划界不要把 VAE tiling、patch 并行、compile 或 offload 称作 VAE 解耦独立的文本编码器 TP 组也不是独立 stage。如果目标版本没有可复用的通用 stage/connector就产出一份可行性结果或 RFC把能力标为not tested而不是发明一个模型私有的生产架构。性能验收清单对每个被接受的优化要求全部满足前后使用同一固定 workload 与同一正确性产物实际快路径的执行由日志/trace/计数器证明同时报告局部指标与端到端延迟/吞吐/内存显式 warmup 后至少三次原始固定工作重复或足够覆盖所有声明尾分位点的到达负载样本集声明命名硬件、后端、拓扑、dtype 与任务范围对不支持的形状/平台有 fallback/rejection 测试abort/错误清理或长时间运行的内存趋势无回归。最后只做窄幅更新特性矩阵不要因为存在某个模型专属 PR就把一个 profile 驱动的方向升级成通用的supported特性。这一原则与 SKILL.md 中定义的证据契约validated/limited/unsupported/not tested四态、逐行记录 checkpoint revision、命令、请求哈希与原始产物完全一致未经验证的 task/backend/cache/quant/offload/topology/hardware 组合一律保持not tested绝不呈现为已支持。总结vLLM-Omni 的性能优化路径是一条先测量、后优化、再验收的纪律性流程测量循环冻结工作负载与指标基线BF16 注意力通过共享角色感知Attention层统一接线并按角色独立资质后端稀疏注意力只在稠密 parity 之后启用并逐任务逐平台验证融合优先模式围绕fused_qk_norm_rope、QKVParallelLinear、SiluAndMul、AdaLayerNorm等共享算子展开保留显式守卫与原生 fallback冗余计算通过提升请求不变量、消除逐层同步与无用 mask 来消除USP 优化关注 all-to-all 边界、GQA padding 与加速传输资质文本编码器与 VAE 解耦以五问五测为前提最终一切以端到端证据和严格验收清单为准。这套方法论既适用于 MiniMax-H3、Boogu-Image 等已在仓库中落地的模型也适用于任何即将加入 vLLM-Omni 的新扩散模型。【免费下载链接】vllm-omniA framework for efficient model inference with omni-modality models项目地址: https://gitcode.com/GitHub_Trending/vl/vllm-omni创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 vLLM-Omni 扩散模型性能模式实战：从测量循环到算子融合、USP 与解耦的完整优化指南
vLLM-Omni 扩散模型性能模式实战从测量循环到算子融合、USP 与解耦的完整优化指南【免费下载链接】vllm-omniA framework for efficient model inference with omni-modality models项目地址: https://gitcode.com/GitHub_Trending/vl/vllm-omni本指南以 vLLM-Omni 仓库中production-add-diffusion-model技能的performance-patterns.md为骨架系统讲解把一个 Day-0 可跑的扩散模型DiT推向可测量、可验收、可上线的过程中必须掌握的性能方法论与代码模式先讲如何建立正确的测量循环再逐一拆解更快的 BF16 注意力、稀疏注意力、融合优先的算子模式Q/K RMSNormRoPE、融合 QKV、Gate-upSwiGLU、AdaLN、冗余计算消除、USP 序列并行通信优化、文本编码器与 VAE 解耦最后给出严格的性能验收清单。读完本文你将掌握一套先固定工作负载、一次只做一个优化、以端到端指标为准的工程化优化流程并能在仓库源码中定位每一类优化对应的共享算子与实现证据。测量循环一切优化都从固定工作负载开始性能文档的第一条铁律是不要以优化清单开头要以固定的正确性与性能工作负载开头。优化前必须先冻结一个可复现的 workload 描述它至少应包含model/checkpoint revision prompt/media hashes seed dimensions/frames scheduler/sigma/steps/guidance task dtype backend topology hardware只有模型/权重版本、输入哈希、随机种子、尺寸/帧数、调度器/sigma/步数/引导强度、任务类型、dtype、注意力后端、并行拓扑、硬件全部固定才能保证前后对比有意义。需要捕获的指标族阶段耗时stage timeload、encode、conditioning、denoise、VAE decode/postprocess 各阶段时间GPU tracekernel 时长与个数、launch 间隔、显存分配、H2D 拷贝、collective集合通信开销CPU trace预处理、逐层同步、Python dispatch、临时文件内存每 rank 的 HBM 峰值、进程树的 host PSS输出链路device 准备、D2H/IPC、payload 字节数、codec 墙钟与进程 CPU、事件循环等待、客户端物化、以及 signed residual不能被掩盖到其他阶段的剩余耗时服务侧队列时间、吞吐、错误、请求混合、分布采样。A/B 方法论对于固定工作的单请求 A/B先跑一次显式 warmup 策略再做至少三次测量重复保留每一次原始 run报告 median/mean 并附范围。不要从三次样本推导 p95/p99尾延迟必须声明到达模型arrival model并采集足够多的成功请求才能支撑所声称的分位点与置信度。一次只实现一个优化重跑 parity同时报告局部 kernel 与端到端两方面的 delta——一个 kernel 赢了但端到端指标毫无变化不能作为生产速度声明。文档还强调优先使用目标版本已暴露的 pipeline profilerPyTorch/CUDA/ROCm/NPU/XPU 各自使用平台配套的 profiling 工具不要把每个平台都强行塞进 CUDA-only 的工作流。更快的 BF16 注意力共享角色感知 Attention 层与后端资质共享注意力接线Shared attention wiring不要每个模型自研一套注意力实现而是使用vllm_omni.diffusion.attention.layer中共享的Attention层并传入真实的局部TP 切分后head 数与 QKV 布局。标准接线如下import torch from vllm.model_executor.layers.linear import ( QKVParallelLinear, RowParallelLinear, ) from vllm_omni.diffusion.attention.layer import Attention class SelfAttention(torch.nn.Module): def __init__( self, hidden_size, head_dim, num_heads, num_kv_heads, quant_configNone, prefix, ): super().__init__() self.head_dim head_dim self.qkv_proj QKVParallelLinear( hidden_sizehidden_size, head_sizehead_dim, total_num_headsnum_heads, total_num_kv_headsnum_kv_heads, biasFalse, quant_configquant_config, prefixf{prefix}.qkv_proj, ) self.attention Attention( num_headsself.qkv_proj.num_heads, num_kv_headsself.qkv_proj.num_kv_heads, head_sizehead_dim, softmax_scalehead_dim**-0.5, causalFalse, qkv_layoutBSND, roleself, prefixprefix, ) self.out_proj RowParallelLinear( num_heads * head_dim, hidden_size, biasFalse, input_is_parallelTrue, quant_configquant_config, prefixf{prefix}.out_proj, ) def forward(self, x, attn_metadataNone): qkv, _ self.qkv_proj(x) q_size self.qkv_proj.num_heads * self.head_dim kv_size self.qkv_proj.num_kv_heads * self.head_dim q, k, v qkv.split([q_size, kv_size, kv_size], dim-1) q q.unflatten(-1, (self.qkv_proj.num_heads, self.head_dim)) k k.unflatten(-1, (self.qkv_proj.num_kv_heads, self.head_dim)) v v.unflatten(-1, (self.qkv_proj.num_kv_heads, self.head_dim)) out self.attention(q, k, v, attn_metadata) out out.flatten(-2) out, _ self.out_proj(out) return out关键点在调用Attention之前按官方顺序先加 Q/K norm 与 RoPE见下一节融合模式保持融合 QKV 与输出投影的 checkpoint 映射完整Attention层内部会依据role与全局 diffusion 配置解析后端。从 layer.py 的实现可以看到它支持role、role_category、qkv_layout、scatter_idx/gather_idxUlysses 注意力、skip_sequence_parallel跳过序列并行通信、custom_attention模型自有 kernel、allow_fp32_fallbackFP32 输入自动回退 SDPA等参数并且当ring_degree 1时会自动挂载RingParallelAttention。后端资质Backend qualification共享Attention类并不自动保证每个后端对 packed 样本都安全或更快。对每个角色self、cross-attention 类别、token refiner 等都要独立记录实际选中的后端从日志/trace 中确认支持的 dtype、head size、QKV 布局、mask/metadata、packed varlen 行为平台与计算能力TP/SP/Ring/AllGather-KV 限制fallback 行为以及 fallback 是否会改变声称的性能。验证时要使用长度不等的 packed 样本与 dense SDPA 对比核对cu_seqlens、最大序列长度、padding 排除、模态边界、CFG 归属与输出切片。不要传了一个attn_mask就假设它被消费了——要检查后端契约与 trace。从 selector.py 可以看到后端的解析优先级是attention_config.per_role[role]→per_role[role_category]→attention_config.default→ 平台默认。TRTLLM、Flash Attention、cuDNN、NPU varlen 等后端是相互独立的资质行Ring 与 hybrid 路径需要各自的边界 parityAllGather-KV 与特定后端组合可能直接 fail fast例如 layer.py 中TRTLLM_ATTN与allgather_degree 1组合会抛错这些行为要写进文档保持可见。稀疏注意力是后端/模型/平台优化不是低精度稠密注意力的同义词稀疏注意力只有在稠密 BF16 parity 达成之后才能开启。对每个候选任务/形状/拓扑必须声明稀疏策略、选定的后端/kernel、block/head 维度、支持的硬件、外部 kernel/包版本与 fallback证明稀疏 kernel 真的执行了并报告实际稀疏率——一个 enabled 标志不够用同种子对比逐 block 输出、denoise 轨迹与最终产物与稠密注意力增加模态专属指标视频的时序一致性与音画同步、音频的频谱/音质、以及相关的 prompt/图像指标报告注意力时间与完整端到端时间包含预处理/index-build 开销、峰值内存与精确的稠密对照组逐任务、逐形状验证不要把一个 T2V 结果外推到 I2V/Ref2V也不要把 NPU 的稀疏支持外推到 CUDA/ROCm/XPU。如果 Ring 绕过了选定后端或某个稀疏设置会被静默忽略就应拒绝该组合。只有当 recipe 把该行标注为limited且不声称稀疏加速时才允许保留可见的稠密 fallback。融合优先的算子模式优先共享算子 显式守卫 原生 fallback融合的准则是优先使用当前仓库的共享/自定义算子在算子、block、轨迹、产物四个层面证明 parity再用 trace 证明 launch 次数/字节数更少或延迟更低。Q/K RMSNorm RoPE先给出通用跨平台基线两个独立分派算子不是自动一个融合 kernelfrom vllm_omni.diffusion.layers.norm import RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk self.norm_q RMSNorm(head_dim, epseps) self.norm_k RMSNorm(head_dim, epseps) # Example only: set these booleans from the official checkpoint contract. use_neox_layout False cos_sin_are_half_dim True self.rope RotaryEmbedding( is_neox_styleuse_neox_layout, half_head_dimcos_sin_are_half_dim, ) q self.norm_q(q) k self.norm_k(k) q, k apply_rope_to_qk(self.rope, q, k, (cos, sin))注意RotaryEmbedding(is_neox_style...)的语义False表示 interleaved/GPT-J 风格偶数奇数维配对True表示 NeoX 风格前半/后半配对不要假设默认是 NeoX。真实代码中要把示例布尔值替换为官方 checkpoint 契约。部分旋转维partial rotary通过旋转前缀 直通后缀拼接实现如果 3D cos/sin 含 batch 轴当前共享路径假设值在 batch 上共享逐样本不同位置需要经过验证的表示/路径。对 packed 的[tokens, heads, head_dim]张量配合非 interleaved 的[cos | sin]表使用共享融合边界from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope q, k fused_qk_norm_rope( q, k, self.norm_q.weight, self.norm_k.weight, rope_table, self.norm_q.variance_epsilon, )从 fused_qk_norm_rope.py 的实现可以看到CUDA 快路径当前特化 BF16head_dim128、rotary_dim96MiniMax-H3 几何契约interleaved 模式支持任意偶数rotary_dim head_dim 256其他合法输入走 eager 实现_apply_rope_table支持 interleaved 与 half-split 两种配对语义分别对应 Boogu-Image 与 MiniMax-H3AscendNPU路径通过mindiesd.rotary_position_embedding/torch_npu.npu_rotary_mul组合实现同样限定 MiniMax-H3 几何函数以自定义算子注册torch.ops.vllm_omni.fused_qk_norm_rope并且提供VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS环境变量token 数低于交叉点时 fused 路径是 host-bound 的反而更慢Boogu-Image 在单张 H200 上实测默认 2048 tokens0表示总是融合。融合 QKV 投影使用上文的QKVParallelLinear。资质要点num_heads与num_kv_heads必须能被 TP 整除使用局部TP 切分后split 尺寸TP 之后绝不复用全局计数在交给 vLLM loader 前完成 Q/K/V checkpoint 顺序转换所有 Q、K、V 分片都必须存在严格加载缺失即失败启用 FP8 等方法时核对量化 scale/loader 映射。Gate-up 投影 融合 SwiGLUfrom vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, RowParallelLinear, ) self.gate_up_proj MergedColumnParallelLinear( input_sizehidden_size, output_sizes[intermediate_size, intermediate_size], biasFalse, quant_configquant_config, prefixf{prefix}.gate_up_proj, ) self.act_fn SiluAndMul() self.down_proj RowParallelLinear( input_sizeintermediate_size, output_sizehidden_size, biasFalse, input_is_parallelTrue, quant_configquant_config, prefixf{prefix}.down_proj, ) gate_up, _ self.gate_up_proj(x) x self.act_fn(gate_up) x, _ self.down_proj(x)checkpoint 的gate_proj与up_proj两个分片都要映射进 merged 参数且两者都要求存在。文档特别提醒merged 投影后跟F.silu(gate) * up与SiluAndMul融合激活 kernel不是同一个声明——要通过 trace 确认实际用的是哪一个。AdaLN只有当共享层的数学/返回契约与官方模型一致时才使用共享层from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm self.adaln AdaLayerNorm( hidden_size, elementwise_affineFalse, eps1e-6, ) x self.adaln(x, scale[:, None, :], shift[:, None, :])AdaLayerNormZero、AdaLayerNormZeroSingle及连续变体包含不同的投影、广播与 tuple 输出不能互相 drop-in 替换。从 adalayernorm.py 可以看到AdaLayerNorm在 NPU 上若有mindiesd会走layernorm_scale_shift融合原语否则 fallback 到torch_npu.npu_layer_norm_eval而 CUDA/HIP/XPU 走forward_nativelayernorm(x) * (1 scale) shift——因此它是一个带 NPU 快路径的共享层但不要称它为通用的融合 AdaLN kernel。对时间步集合有限的 sigma schedule可以先 profile 一下每个请求预计算一次 timestep/AdaLN 投影是否值得缓存 key 必须包含精确的 schedule/config/device/dtype 且保持有界。布局、残差与其他候选profiling 之后继续检查bias/dropout/add 或 gated residual 链normlinear 或 normmodulation 边界注意力周围的重复 transpose/reshape/contiguousVAE 的 normalization/activation/convolution 边界真正使用 MoE 的模型的 MoE routing 与 fused experts。融合 residual-plus-norm 或 modulation 链时保留未融合的物化materialization边界如果 residual 在 RMSNorm 前以 BF16 存储就要在融合 kernel 内先舍入到 BF16 再提升回计算精度直接使用未舍入的 FP32 临时量即使看起来对也会改变 denoise 轨迹。DLO 的物化/释放 hooks、进程组集合通信、请求缓存迁移与清理等生命周期操作要放在编译/融合区域之外除非生命周期安全性被显式测试覆盖。只编译稳定区域并保留 eager fallback。消除冗余计算提升请求不变量Hoist request invariants审计 denoise 循环中只依赖请求、不依赖当前时间步的工作把它们提升到循环之前# Before the denoise loop. condition self.condition_proj(text_embeddings) refined_condition self.token_refiner(condition, packed_metadata) rope_tables self.build_rope_tables(position_ids) attn_metadata self.build_attention_metadata(layout) for step, sigma in enumerate(sigmas): timestep_state self.prepare_timestep_state(sigma) noise self.transformer( latents, refined_condition, rope_tablesrope_tables, attn_metadataattn_metadata, timestep_statetimestep_state, ) latents self.scheduler_step(noise, sigma, latents)候选包括prompt conditioning、token-refiner 输出、reference 编码/扫描、position IDs/RoPE 表、packed masks/metadata、静态 CFG 分支、有限 schedule 的 AdaLN 投影。先证明它们对所有官方任务、LoRA/adapters、prompt-update 特性、形状与 schedule 都不变在连续 batching 中请求专属结果要存在请求状态上而不是 pipeline 单例上。消除逐层同步避免在每层内部做 device-to-host 标量提取如.item()# Bad: forces a GPU/CPU synchronization in every attention layer. max_len int(cu_seqlens[-1].item()) # Better: compute validated host metadata once during request packing and pass it. packed PackedLayout( cu_seqlenscu_seqlens, max_seqlenmax(sample_lengths), ) for block in self.blocks: hidden block(hidden, packed_layoutpacked)用 trace 验证同步确实消失。不要为了去掉.item()而把正确的动态值换成过期的常量。移除未使用的 mask 与分配如果 packed 注意力消费cu_seqlens而忽略 dense mask就不要构造 dense mask。安全时复用不可变请求元数据与预分配缓冲。优先用直接索引的 scatter/gather 替代构造大型临时张量——前提是重复索引语义与梯度推理时通常禁用保持正确。任何跨请求缓存都需要有界生命周期且 key 必须完整覆盖 shape、layout、device、dtype、schedule、task、adapter/LoRA 以及任何会改变张量的值否则保持请求级作用域。VAE 与 reference 工作profile 重复的 VAE decode/postprocess、重复的 reference 解码/扫描、不必要的 host-device 往返与整帧物化。启用 VAE tiling、slicing、patch 并行或按需暂存时必须配套 seam、顺序、颜色/范围与内存 parity 测试。USP 通信优化从能跑到通信被优化USP 能跑不等于 USP 通信被优化了。对 Ulysses、Ring、hybrid 与 AllGather-KV 要分开profilecollective 的数量、字节数、时长、所在 stream 以及与计算的 overlapQ/K/V 与输出的 all-to-all 边界pack/unpack、padding、metadata 广播、transpose 与 contiguous 拷贝head 整除性或高级 uneven-head 支持同一 workload/互连上从 1 到 N 设备的扩展效率长度不等 packed 样本的负载均衡。_sp_plan的 split/gather 边界应声明在有意义的模块输出处而不是任意的 Python 语句packed metadata 要按与 token 相同的 ownership 模型做 shard/gather。加速 Ulysses 传输--ulysses-a2a-permute当目标版本提供加速的 Ulysses 传输如--ulysses-a2a-permute时要把它与普通 Ulysses分开资质。从 ulysses.py 的实现可见加速路径的启用条件是严格 scatter-heads/gather-sequence 布局scatter_idx 2 and gather_idx 1配置项ulysses_a2a_permute定义于 omni_config.py。实践要点证明严格布局确实选中了快路径记录一次性 JIT/readiness 开销在 CUDA graph 捕获前对最大 workspace 形状做 prewarm让 grow-only workspace 保持在单条 stream 上在进程组销毁前执行 cleanup。对 GQA验证Q_heads % KV_heads 0先 pad KV heads再由原始比例推导 Q 的 padding。该文件还内置了一个警告机制当 advanced_uaa 的 GQA padding 把 Q head 数膨胀超过 1.5 倍时最坏情况 MQA 从 H 膨胀到 N*H会一次性告警提示用户选择更合适的ulysses_degree。同时要同时测试 contiguous 与 row-strided 两种 QKV 布局快拷贝路径必须保持逐位一致与 packed 边界。复制型 cross-attention 的 skip_sequence_parallelcross-attention 有时使用复制的文本 K/V只有视频/音频 Q 做序列 shard。这种特定场景下skip_sequence_parallelTrue可以避免错误或浪费的 Q/K/V 重分布self.cross_attention Attention( num_headslocal_heads, num_kv_headslocal_kv_heads, head_sizehead_dim, softmax_scalehead_dim**-0.5, causalFalse, rolecross, skip_sequence_parallelTrue, prefixf{prefix}.cross_attn, )只有在证明 K/V 确为复制、Q/输出 ownership 正确之后才使用——否则它会禁用必要的通信。报告该优化是否减少了 collective 字节/时间并改善端到端扩展。文本编码器与 VAE 解耦先测量再架构解耦disaggregation是一个需要实测的架构决策不是默认方案。先收集五问五测问题需要测量的量阶段是否足够大目标并发与形状下的 encode/decode 占比传输是否划算tensor/media 字节、序列化、connector 延迟与带宽是否有复用prompt/reference 复用率、可缓存性、fan-out能否独立扩展各阶段的队列饱和度与资源利用率契约是否稳定shape、dtype、mask/tag/position 元数据、顺序文本编码器候选定义显式 tensor 契约hidden states、masks、tags 或 position IDs、dtype、shape、task/checkpoint revision、request ID。支持禁用本地编码器避免 DiT 意外重算它。测试所有任务、prompt 长度、部分 TP 组成员、重试、背压、abort、connector 清理与独立副本扩展。VAE 候选如果 encode 与 decode 契约都存在要分开定义latent shape、scale/shift、dtype/range、tiling/chunk 顺序、输出媒体顺序与阶段元数据。用 latent/media 传输量对比本地 decode 成本。测试 tiling 接缝、相关场景的 RNG/确定性、chunk 顺序、codec/postprocess、背压、abort 与 worker 失败。文档明确划界不要把 VAE tiling、patch 并行、compile 或 offload 称作 VAE 解耦独立的文本编码器 TP 组也不是独立 stage。如果目标版本没有可复用的通用 stage/connector就产出一份可行性结果或 RFC把能力标为not tested而不是发明一个模型私有的生产架构。性能验收清单对每个被接受的优化要求全部满足前后使用同一固定 workload 与同一正确性产物实际快路径的执行由日志/trace/计数器证明同时报告局部指标与端到端延迟/吞吐/内存显式 warmup 后至少三次原始固定工作重复或足够覆盖所有声明尾分位点的到达负载样本集声明命名硬件、后端、拓扑、dtype 与任务范围对不支持的形状/平台有 fallback/rejection 测试abort/错误清理或长时间运行的内存趋势无回归。最后只做窄幅更新特性矩阵不要因为存在某个模型专属 PR就把一个 profile 驱动的方向升级成通用的supported特性。这一原则与 SKILL.md 中定义的证据契约validated/limited/unsupported/not tested四态、逐行记录 checkpoint revision、命令、请求哈希与原始产物完全一致未经验证的 task/backend/cache/quant/offload/topology/hardware 组合一律保持not tested绝不呈现为已支持。总结vLLM-Omni 的性能优化路径是一条先测量、后优化、再验收的纪律性流程测量循环冻结工作负载与指标基线BF16 注意力通过共享角色感知Attention层统一接线并按角色独立资质后端稀疏注意力只在稠密 parity 之后启用并逐任务逐平台验证融合优先模式围绕fused_qk_norm_rope、QKVParallelLinear、SiluAndMul、AdaLayerNorm等共享算子展开保留显式守卫与原生 fallback冗余计算通过提升请求不变量、消除逐层同步与无用 mask 来消除USP 优化关注 all-to-all 边界、GQA padding 与加速传输资质文本编码器与 VAE 解耦以五问五测为前提最终一切以端到端证据和严格验收清单为准。这套方法论既适用于 MiniMax-H3、Boogu-Image 等已在仓库中落地的模型也适用于任何即将加入 vLLM-Omni 的新扩散模型。【免费下载链接】vllm-omniA framework for efficient model inference with omni-modality models项目地址: https://gitcode.com/GitHub_Trending/vl/vllm-omni创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考