Repository navigation
[Core] Introduce a generic cross-attention abstraction with LingBot integration - #8859
Conversation
…V lifecycle management Signed-off-by: tsinghua-code <zr-liu25@mails.tsinghua.edu.cn>
Signed-off-by: tsinghua-code <zr-liu25@mails.tsinghua.edu.cn>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
|
This PR appears to belong to: docs/design/module/diffusion/index.md. Module owners: @wtomin @david6666666 @Isotr0py Routing: @wtomin via module of the changed files, CODEOWNERS; @david6666666 via module of the changed files; @Isotr0py via module of the changed files @tsinghua-code, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer. Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment. |
Self-review summary (with Codex assistance):
|
Signed-off-by: tsinghua-code <zr-liu25@mails.tsinghua.edu.cn>
Signed-off-by: tsinghua-code <zr-liu25@mails.tsinghua.edu.cn>
|
I suggest keeping the execution contract, conditioning-cache lifecycle, and first-block sharing changes in this PR, with clear internal boundaries and separate validation. Before settling the public contract, we should use other models to check which assumptions are specific to LingBot:
Based on these cases, I would refine the design as follows:
For validation, I would prioritize LingBot + Helios as concrete integrations, then add focused conformance coverage for masked cross-attention, GQA, both LTX2 parallel directions, independent text/image softmaxes, packed-request isolation, and selective cache invalidation. We do not need to migrate every model in this PR, but the contract should represent these requirements and clearly distinguish implemented support from extension points. |
Signed-off-by: tsinghua-code <zr-liu25@mails.tsinghua.edu.cn>
|
@tzhouam Thanks for the detailed review on the contract and cache boundaries. I used the Helios, Wan/DreamZero, Stable Audio, LTX2, and H3 cases to refine the shared contract. The follow-up is in production-only commit
Validation summary (old and new use the same checkpoint, prompt, seed, and execution path within each row):
Scope limits: This version isolates requests by BSND batch row or an explicit validated pair mask. Native packed/varlen cross metadata is rejected because its independent Q and K/V boundaries and pair mapping are not yet represented; H3's packed joint self-attention remains on its existing path. DreamZero's selective window-reset cache behavior and Sana model integration are not migrated. The real-checkpoint rerun on the combined worktree was single-GPU; the earlier multi-GPU conformance and model runs are separate evidence. Cross-request first-block sharing remains disabled by default: prior same-version sharing-on/off measurements showed first-chunk latency increases of 6.01% (TP2), 11.59% (SP2), and 21.19% (TP2+SP2), so I am not claiming a speedup. Compilation combined with first-block sharing is still unsupported. Would batch-row isolation plus explicit pair-mask isolation, with native packed cross rejected, satisfy the packed-boundary requirement for this PR? I can keep native varlen cross as a separately specified extension if so. AI assistance was used for implementation, validation, and preparing this response. |
Signed-off-by: tsinghua-code <zr-liu25@mails.tsinghua.edu.cn>
Head branch was pushed to by a user without write access
Omni ReviewBot: CI is red on this head@tsinghua-code required checks failed on Please fix the failure or rerun CI; this note is updated in place when the failing checks change, the head goes green, or it moves. |
Resolve the SDPAImpl._forward_impl conflict with vllm-project#8859 by moving its FP32-bias handling into the shared SDPA helpers: _sdpa_uses_expanded_kv expands K/V when a floating bias differs from Q's dtype, and _sdpa_attention runs that case in the math kernel. Eager calls and the compiled GQA custom op therefore both keep the fix, and the K/V-expansion debug log stays skipped while compiling. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: princepride <wangzhipeng628@gmail.com>
…etirement hook Main's runner now calls pipeline.release_step_request(state) for every retired step state (vllm-project#8859). SenseNova implements it: drop the request's step context, including a prepare-phase decode cursor, and clear the flash KV buffers of its denoise caches when an abort or failure retires the request before post_decode. Requests that retire inside the resumable prepare phase, after a prepare failure or with a prepare-only output, now go through release_step_requests as well, so the pipeline sees the same callback as for a denoised request. Co-authored-by: zijianc2 <157244773+zijianc2@users.noreply.github.com> Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>
…b4-step-batching Conflict resolutions and review-feedback adaptations for main through 841cb06 (vLLM v0.32.0 alignment, vllm-project#8254, vllm-project#8268, vllm-project#8380, vllm-project#8859, vllm-project#6285): - diffusion_model_runner: keep our release_step_requests (returns int, invalidates the input batch) and fold in main's vllm-project#8859 pipeline retirement hook, so every popped state reaches pipeline.release_step_request on finish, abort, and failure. - sensenova_u1 pipeline: add SenseNovaU1Pipeline.release_step_request delegating to release_step_state and drop the weakref backstop now that retirement is deterministic; hoist the CFG interval decision, vision grid, and noise-scale embedding into _denoise_context per vllm-project#8254, computed once in prepare_encode (step path) or at loop entry and passed through _denoise_one into the now use_cfg-required _denoise; read od_config.step_execution inside pre_process_func so streaming_output text requests get their prepare budget. - tests: resolve admission pre-processing through get_diffusion_pre_process_func (pins the registry entry), drop the bare pytest.mark.cuda disallowed by vllm-project#8380's check_test_marks, add a denoise_step CFG-boundary test on the real _denoise, and adapt stubs to the set_request_adapter rename from vllm-project#8268. SenseNova suite: 201 passed. Engine/scheduler/runner/step/resumable core files: 568 passed, 3 skipped. Signed-off-by: shihongzhi <shi65881583@gmail.com>
Contributor intent
1. What outcome do you want to achieve?
设计并实现通用 cross-attention 执行抽象,以 LingBot World 作为首个接入实例,验证去噪输出、latent 和视频帧与原实现一致,并管理请求内及受约束的跨请求 K/V 生命周期。
2. What is the highest-leverage variable for that outcome?
建立清晰、可复用的 cross-attention 执行与缓存契约,并保证模型接入后计算结果与原实现一致,缓存复用不会造成请求间状态污染。
3. How will this work move that variable?
统一 cross 计算语义及执行契约,接入请求内缓存有效性和生命周期管理,并为跨请求首块视频 K/V 共享定义复用及隔离条件,通过成对计算结果和性能测量进行验收。
4. Whose customer or user needs does this address, and what are those needs?
主要面向 vLLM-Omni 的模型接入开发者和推理服务维护者。他们需要复用统一的 cross-attention 执行及缓存管理能力,减少模型专用代码,同时保证计算正确、请求隔离和资源生命周期可控。
1. 背景与目标
关联 RFC #8659、PR #8064。
本 PR 设计并实现通用 cross-attention 执行抽象,明确请求内 K/V 缓存及受约束的跨请求 K/V 共享的有效性、隔离与生命周期。LingBot World 是首个接入实例;Helios 的普通 tensor conditioning 缓存也接入公共生命周期规则。
包含三个部分:
共享执行抽象和 LingBot 参考接入属于 RFC #8659 的原始范围;缓存生命周期及跨请求首块共享是本 PR 的范围扩展。
2. 设计与职责边界
通用执行层
模型继续负责 Q/K/V/O 投影、normalization、位置编码及条件组合。通用执行层消费已经投影的 Q/K/V,负责执行契约检查、语义传递及后端执行。
role用于配置选择,计算语义通过显式computation声明。computation、并行放置与缓存策略分别声明;cross 语义本身不会强制跳过 SP,也不会自动启用 K/V 缓存。Attention → AttentionBackend → AttentionImpl架构。缓存边界
相同首帧图片不足以证明首块 K/V 可以共享;文本、随机 latent、相机条件等同样影响其值。
3. 实现内容
3.1 通用 cross-attention 执行抽象
role或缓存策略。role == "cross"判断。3.2 请求内 conditioning K/V 生命周期
将缓存有效性和生命周期从 AR 分页存储中解耦:公共
ConditioningKVScope统一查找、完整发布、失效和关闭,ConditioningKVStorage定义存储接口。LingBot 继续使用现有 AR 存储;Helios 通过TensorConditioningKVStorage使用普通 tensor 存储,原 AR 导入路径仍兼容。缓存有效性不只依赖 tensor 地址和 shape。Helios 检查源 tensor 对象及可用的修改版本;无法可靠取得版本时使用快照判断,保留 #8064 针对地址复用的防护。普通路径在请求退出时关闭缓存;stepwise 在完成、失败和释放时清理。
3.3 受约束的跨请求首块 K/V 共享
本 PR 的新增提交
6f55ecc2f包含 18 个生产代码文件:diffusion/attention/{layer.py,capabilities.py,backends/{abstract.py,flash_attn.py,sdpa.py,utils/attn_runtime_selector.py},parallel/{cross.py,factory.py}}diffusion/attention/conditioning.py、experimental/ar_diffusion/kv_cache/conditioning.pydiffusion/models/{lingbot_world/pipeline.py,helios/{conditioning_cache.py,cross_attn_cache.py,helios_transformer.py,pipeline_helios.py},stable_audio/stable_audio_transformer.py,wan2_2/wan2_2_transformer.py}diffusion/worker/diffusion_model_runner.py上述生产代码路径均位于
vllm_omni/。另包含两个既有 CI 测试文件的接口适配:test_pipeline_lingbot_world.py和test_paged_attention_adapter.py。独立验证脚本、benchmark、生成示例及视频产物不进入代码改动。生成效果:大本钟视频对照
相同 checkpoint、首帧、prompt、seed 和生成配置,每版生成 21 chunk、249 帧,约 15.56 秒。以下视频为模型生成结果,首帧为 AI 生成素材,不是实拍或独立正确性证据。
旧版:模型原有 cross-attention
查看或下载旧版视频
新版:通用 cross-attention
查看或下载通用 cross-attention 视频
素材与生成条件
robbyant/lingbot-world-v2-14b-causal-fast-diffusers59cccf49f2d2dd27418ae7a04b82b10868d455c2forward下载原始 PNG
首帧原始文件 SHA-256:
复现命令
安装对应版本的 vLLM-Omni,准备上述模型快照与原始 PNG。在两个终端分别设置以下变量,将路径替换为本机实际路径:
终端一启动服务:
终端二在服务就绪后生成视频:
每条请求生成完整的 21 chunk,保存最后一条请求的视频。
--warmup-chunks 6仅影响性能统计。对旧版执行相同命令时,将
LINGBOT_REPO指向旧版 checkout,并将输出文件名改为big-ben-baseline。两版服务依次启动,使用同一张 GPU、模型快照及输入素材。展示视频来自旧版
cf06215e3与第一版通用实现13c45cf43。第三版开启共享的两次 21-chunk 回归也记录了相同视频哈希及逐阶段零差异。展示视频不标记为当前整理提交重新生成的产物。4. 正确性验收
4.1 被测版本与判定标准
以下 4.2–4.4 表格是较早的第三版生产代码
1d9ca7f1c的验收记录;整理提交93a6c6dbc的定向 CPU 回归为 39 passed, 17 warnings,当时的完整 pre-commit 检查通过。这些历史多卡结果并非新增提交6f55ecc2f上重新运行。新增提交的 LingBot 与 Helios 联合单卡回归见 4.5。以下历史 LingBot 验收的软件环境:PyTorch 2.13.0、vLLM 0.31.0、Transformers 5.14.1、Diffusers 0.40.0、NumPy 2.3.5。
4.2–4.4 的 LingBot 真实模型对照使用 BF16、832×480、4 步去噪;每条完整请求生成 10 chunk、117 帧。每组成对请求检查 40 个去噪输出、40 个 latent 输出、全部 117 帧的逐元素一致性及数组文件 SHA-256。多卡按相同 TP/SP 拓扑对照,不要求不同拓扑之间逐位一致。
4.2 计算等价性结果
三种模式为旧实现、新实现关闭共享、新实现开启共享。成对比较包含旧/新、共享开/关和共享冷/暖,同一请求可参与多个比较。单卡输入序列:原输入 → 原输入 → 换文本 → 换 seed → 换相机 → 换图片 → 恢复原输入。编译行仅表示关闭共享的比较通过。
4.3 缓存生命周期与隔离
LRU、替换及交错测试包含在 82 项 CPU 回归中。
4.4 四卡取消与资源增长
配置:TP2+SP2 eager、开启首块共享、真实 checkpoint。
stopped=true资源增长比较第 5–8 条请求与最后 4 条请求的占用中位数。本压力组检查最终媒体输出,中间阶段由独立数值测试覆盖。请求可重叠到达,GPU 执行仍串行。真实接口未强制复用同一请求 ID,该场景由 CPU 测试覆盖。
4.5 当前提交的 LingBot 与 Helios 联合回归
生产代码
6f55ecc2f上,联合工作区的 99 项定向测试通过,0 失败、0 跳过。每一行均在相同 checkpoint、输入、seed 和执行路径下比较旧版与新版,普通与 stepwise 分别成对比较。测试脚本及原始报告保存在服务器验证目录中,未加入本 PR。Helios 旧版已显式使用 cross 计算语义,因此其成对结果验证本次缓存接入。联合真实 checkpoint 重跑为单卡;此前 SP2、TP2+SP2 的执行契约测试和 Helios 多卡模型测试是独立版本的证据,不等同于当前联合代码的多卡重跑。
5. 性能结果
5.1 旧版专有路径与新版通用路径
此前完成的 eager 基准,两侧均关闭跨请求首块共享。每版运行 4 条请求,排除首条暖机,报告 3 条暖请求的中位数。新路径包含通用执行契约及同期 conditioning K/V 生命周期变更,端到端差异不能全部单独归因于 attention 入口。
每个 stepwise 配置的 8 条请求、80 个媒体 chunk 全部一致;普通路径 8 条请求的最终帧数组哈希一致。以上配置的三个指标均未出现超过 3% 工程容差的退化;这不等于严格零退化,也不是统计等效性证明。
5.2 同一新版的跨请求共享开关对照
使用关闭 → 开启 → 开启 → 关闭的顺序。每次启动运行 4 条请求,排除首条冷请求,每种状态累计 6 条暖请求。计时期间关闭中间 tensor 记录。
每种拓扑的 16 条请求、160 个媒体 chunk 输出均一致。
RSS 为服务进程合计占用中位数,包含 CPU 快照及其他内存,不代表峰值。四卡开启组包含一条 9248.87 ms 的首 chunk 高延迟样本,原样保留。测量在共享主机进行,样本量有限。
跨请求共享的输出正确性通过,但三个拓扑的首 chunk 性能均未达到无退化要求,也超过 3% 容差。 当前共享默认关闭,不宣称已取得性能收益。
5.3 Helios 缓存接入的单卡性能
真实 checkpoint、固定 10-chunk 配置的旧版/新版/旧版独立重启对照:旧版 6 条请求中位 40.0020 s,新版 3 条请求中位 39.7288 s,观测差异 −0.68%。样本量有限,不能据此证明普遍加速;该数据也不能归因于 LingBot 通用执行路径。
6. 支持范围与已知限制
max_num_seqs > 1被拒绝。通用执行契约可按 BSND batch 行或显式 pair mask 隔离请求,但不等于 LingBot 已支持多请求 AR 批处理。AI assistance: 使用 Codex 辅助实现、测试执行、代码整理及 PR 正文撰写。生成视频与首帧已在上文标注。