Skip to content

[Core] Introduce a generic cross-attention abstraction with LingBot integration - #8859

Merged
tzhouam merged 7 commits into
vllm-project:mainfrom
tsinghua-code:codex/track-a-8659-production-only
Oct 17, 2026
Merged

tzhouam merged 7 commits into
vllm-project:mainfrom
tsinghua-code:codex/track-a-8659-production-only

Conversation

@tsinghua-code

@tsinghua-code tsinghua-code commented Oct 13, 2026 •

Copy link
Copy Markdown
Contributor

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 缓存也接入公共生命周期规则。

包含三个部分:

  1. 统一 cross-attention 执行语义与接口。
  2. 管理同一请求内 conditioning K/V 的复用、失效和释放。
  3. 实现受完整输入匹配约束的跨请求首块视频 K/V 共享。

共享执行抽象和 LingBot 参考接入属于 RFC #8659 的原始范围;缓存生命周期及跨请求首块共享是本 PR 的范围扩展。

2. 设计与职责边界

通用执行层

模型继续负责 Q/K/V/O 投影、normalization、位置编码及条件组合。通用执行层消费已经投影的 Q/K/V,负责执行契约检查、语义传递及后端执行。

  • role 用于配置选择,计算语义通过显式 computation 声明。
  • Q/KV 来自独立序列;长度相等不能据此推断为 self-attention。
  • 已实现的 cross 契约支持 BSND、合法的 Q/KV head 几何(含 GQA),以及本地 Q/复制 K/V、复制 Q/序列分片 K/V 两种放置。是否 causal 单独声明。
  • Q 有效性、KV 有效性和 Q–KV 配对限制分别声明;请求按 BSND batch 行或显式配对 mask 隔离。原生 packed/varlen cross metadata 暂不支持。
  • computation、并行放置与缓存策略分别声明;cross 语义本身不会强制跳过 SP,也不会自动启用 K/V 缓存。
  • 复用现有 Attention → AttentionBackend → AttentionImpl 架构。
  • 不支持的组合明确报错,保留已有调用方式的兼容处理。

缓存边界

缓存类型 同一请求内复用 跨请求复用 约束
文本 conditioning cross-attention K/V 支持 禁止 请求实例、层、分支及相关版本有效
首个视频块 self-attention K/V 支持 有条件支持 完整输入、模型依赖及存储表示匹配
请求运行中的可变视频 K/V 支持 禁止直接共享可变存储 每个请求使用独占视频页

相同首帧图片不足以证明首块 K/V 可以共享;文本、随机 latent、相机条件等同样影响其值。

3. 实现内容

3.1 通用 cross-attention 执行抽象

  • 增加明确的 self/cross 计算语义及两种 cross 并行放置声明;计算语义不等同于配置 role 或缓存策略。
  • 在通用入口检查 Q/K/V 布局、长度、合法 GQA head 几何、head dimension、dtype、device、Q/KV/pair mask 及请求边界。
  • 根据实际几何、mask 和并行放置检查后端能力;允许的组合按规则回退 SDPA,不支持的显式组合报错。Mask 内容在准备 metadata 时检查,避免每层重复同步。
  • 将计算语义传递到执行上下文及后端,使 FlashAttention 不再仅依赖 role == "cross" 判断。
  • LingBot 保持原有权重路径、精度、TP/SP 行为及模型计算职责;Stable Audio 使用原生 GQA head 数,Wan2.2 显式声明 cross 语义。模型仍各自执行文本/图像等条件分支并组合结果,LTX2 等模型未整体迁移。

3.2 请求内 conditioning K/V 生命周期

将缓存有效性和生命周期从 AR 分页存储中解耦:公共 ConditioningKVScope 统一查找、完整发布、失效和关闭,ConditioningKVStorage 定义存储接口。LingBot 继续使用现有 AR 存储;Helios 通过 TensorConditioningKVStorage 使用普通 tensor 存储,原 AR 导入路径仍兼容。

  • 同一请求内,条件及投影依赖不变时,可跨 steps 和 chunks 复用。
  • 条件、投影 owner/版本、adapter、逻辑分支或存储表示变化后,旧缓存失效。
  • 构建完成并通过发布检查后,缓存才可被查找。
  • 使用请求实例标识区分复用的外部请求 ID。
  • 处理关闭幂等、构建期间失效及失败清理。
  • 使用 CUDA 就绪事件和消费 stream 记录,保护跨 stream 可见性与在途使用。

缓存有效性不只依赖 tensor 地址和 shape。Helios 检查源 tensor 对象及可用的修改版本;无法可靠取得版本时使用快照判断,保留 #8064 针对地址复用的防护。普通路径在请求退出时关闭缓存;stepwise 在完成、失败和释放时清理。

3.3 受约束的跨请求首块 K/V 共享

  • 共享首个生成视频块的 self-attention K/V。
  • 复用键检查 clean latent、图像条件、文本 embedding、相机条件及相关模型依赖和存储表示。
  • 命中后,将 CPU 快照恢复到新请求独占的 GPU 视频页。
  • 文本 conditioning K/V 始终保持请求隔离。
  • 构建使用代次检查,防止失效后的旧结果重新发布。
  • 对快照容量、替换和已借出对象的生命周期进行约束。
  • 跨请求共享默认关闭。
对象 管理方式
保留会话 复用现有 LRU,淘汰时关闭并释放请求资源
跨请求首块快照 单条 CPU 快照,具有容量限制及替换保护;不是多条共享快照的 LRU

本 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}}
公共缓存与 AR 兼容 diffusion/attention/conditioning.py、experimental/ar_diffusion/kv_cache/conditioning.py
模型接入 diffusion/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}
stepwise 清理 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-diffusers
模型快照 59cccf49f2d2dd27418ae7a04b82b10868d455c2
首帧 随本 PR 提供的大本钟原始 PNG
分辨率 / 精度 832×480 / BF16
执行模式 单卡 eager,关闭跨请求首块共享
seed 42
去噪步数 / flow shift 4 / 5.0
相机动作 forward
输出 21 chunk、249 帧、16 FPS
实测 GPU NVIDIA L20X

大本钟首帧素材

下载原始 PNG

首帧原始文件 SHA-256:

1af3612de7aa9a93994312588f948ccce92bf41da6e6d467d4f89eb985b2d784

复现命令

安装对应版本的 vLLM-Omni,准备上述模型快照与原始 PNG。在两个终端分别设置以下变量,将路径替换为本机实际路径:

export LINGBOT_REPO=/absolute/path/to/vllm-omni
export LINGBOT_MODEL=/absolute/path/to/lingbot-model-snapshot
export LINGBOT_IMAGE=/absolute/path/to/big-ben-input.png
export PYTHONPATH="$LINGBOT_REPO${PYTHONPATH:+:$PYTHONPATH}"

终端一启动服务:

CUDA_VISIBLE_DEVICES=0 vllm serve "$LINGBOT_MODEL" \
  --omni \
  --deploy-config "$LINGBOT_REPO/benchmarks/lingbot_world/configs/single_gpu_eager.yaml" \
  --host 127.0.0.1 --port 18075 \
  --init-timeout 600 --stage-init-timeout 600

终端二在服务就绪后生成视频:

python "$LINGBOT_REPO/benchmarks/lingbot_world/benchmark_lingbot_world_realtime.py" \
  --host 127.0.0.1 --port 18075 \
  --model "$LINGBOT_MODEL" \
  --image "$LINGBOT_IMAGE" \
  --prompt 'Big Ben (the Elizabeth Tower) and the Palace of Westminster in London, England, beside the River Thames. The camera moves slowly forward toward the clock tower in natural daylight.' \
  --num-chunks 21 --camera-pattern forward \
  --width 832 --height 480 --fps 16 --seed 42 --flow-shift 5.0 \
  --sessions 2 --warmup-chunks 6 \
  --save-video /tmp/big-ben-candidate.mp4 \
  --output-json /tmp/big-ben-candidate.json

每条请求生成完整的 21 chunk,保存最后一条请求的视频。--warmup-chunks 6 仅影响性能统计。

对旧版执行相同命令时,将 LINGBOT_REPO 指向旧版 checkout,并将输出文件名改为 big-ben-baseline。两版服务依次启动,使用同一张 GPU、模型快照及输入素材。

比较项 结果
去噪输出 84/84 份完全一致
latent 输出 84/84 份完全一致
视频帧 249/249 帧完全一致
MP4 文件大小 两份均为 11,033,072 字节
MP4 字节及 SHA-256 完全一致

展示视频来自旧版 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 计算等价性结果

路径 完整请求数 成对比较通过数 去噪不一致元素 latent 不一致元素 视频帧不一致元素
单卡普通路径 旧/新各 1 条,共 2 条 1/1 0 0 0
单卡 stepwise eager 三种模式各 7 条,共 21 条 22/22 0 0 0
TP2 eager 三种模式各 2 条,共 6 条 7/7 0 0 0
SP2 eager 三种模式各 2 条,共 6 条 7/7 0 0 0
TP2+SP2 eager 三种模式各 2 条,共 6 条 7/7 0 0 0
TP2+SP2 编译,关闭共享 旧/新各 3 条,共 6 条 3/3 0 0 0

三种模式为旧实现、新实现关闭共享、新实现开启共享。成对比较包含旧/新、共享开/关和共享冷/暖,同一请求可参与多个比较。单卡输入序列:原输入 → 原输入 → 换文本 → 换 seed → 换相机 → 换图片 → 恢复原输入。编译行仅表示关闭共享的比较通过。

4.3 缓存生命周期与隔离

项目 规模 结果
CPU 回归 82 项 82 通过,0 失败,0 跳过
CUDA stream 生命周期 2 项 2 通过,0 跳过
会话 LRU 淘汰及 ID 复用 128 轮 活跃视频页隔离;淘汰释放文本缓存;关闭恢复空闲页
共享首块快照替换 256 次 旧 key 不命中新值;容量受限;已借出快照仍有效
构建与失效/发布交错 64 次 失效后的旧构建不能重新发布
不同输入隔离 分别改变文本、seed、相机、图片 缓存 miss,输出与对应旧版一致
TP2 暖请求命中及恢复 2 个 rank 2/2
SP2 暖请求命中及恢复 2 个 rank 2/2
TP2+SP2 暖请求命中及恢复 4 个 rank 4/4

LRU、替换及交错测试包含在 82 项 CPU 回归中。

4.4 四卡取消与资源增长

配置:TP2+SP2 eager、开启首块共享、真实 checkpoint。

指标 实测结果 判定
完整请求 29 条,每条 10 chunk 媒体输出与对应参考一致
主动取消 6 次 6/6 返回 stopped=true
首 chunk 前取消 3 次 完成取消
首 chunk 后取消 3 次 完成取消
取消与后继请求重叠到达 3 次 后继完整请求输出一致
服务进程合计 RSS 增长 64.91 MiB 低于 256 MiB 门槛
四张卡显存增长 4 / 2 / 0 / 4 MiB 均低于 256 MiB 门槛

资源增长比较第 5–8 条请求与最后 4 条请求的占用中位数。本压力组检查最终媒体输出,中间阶段由独立数值测试覆盖。请求可重叠到达,GPU 执行仍串行。真实接口未强制复用同一请求 ID,该场景由 CPU 测试覆盖。

4.5 当前提交的 LingBot 与 Helios 联合回归

生产代码 6f55ecc2f 上,联合工作区的 99 项定向测试通过,0 失败、0 跳过。每一行均在相同 checkpoint、输入、seed 和执行路径下比较旧版与新版,普通与 stepwise 分别成对比较。测试脚本及原始报告保存在服务器验证目录中,未加入本 PR。

模型与路径 完整请求 去噪输出 latent 输出 视频输出
LingBot 普通 旧/新各 10 chunk 40/40 一致 40/40 一致 117/117 解码帧一致
LingBot stepwise 旧/新各 10 chunk 40/40 一致 40/40 一致 10/10 视频 chunk 及媒体字节一致
Helios 普通 旧/新各 10 chunk 66/66 一致 66/66 scheduler latent、10/10 decode latent 一致 10/10 帧块逐位一致
Helios stepwise 旧/新各 10 chunk 66/66 一致 66/66 scheduler latent、10/10 decode latent 一致 10/10 帧块逐位一致

Helios 旧版已显式使用 cross 计算语义,因此其成对结果验证本次缓存接入。联合真实 checkpoint 重跑为单卡;此前 SP2、TP2+SP2 的执行契约测试和 Helios 多卡模型测试是独立版本的证据,不等同于当前联合代码的多卡重跑。

5. 性能结果

5.1 旧版专有路径与新版通用路径

此前完成的 eager 基准,两侧均关闭跨请求首块共享。每版运行 4 条请求,排除首条暖机,报告 3 条暖请求的中位数。新路径包含通用执行契约及同期 conditioning K/V 生命周期变更,端到端差异不能全部单独归因于 attention 入口。

路径 旧版完整耗时 新版完整耗时 变化 旧版首 chunk 新版首 chunk 变化
单卡普通 26.445 s 26.430 s −0.056% — — —
单卡 stepwise 25.838 s 25.903 s +0.25% 2236.62 ms 2247.58 ms +0.49%
TP2 17.684 s 17.773 s +0.50% 1713.69 ms 1722.47 ms +0.51%
SP2 17.819 s 17.061 s −4.25% 1660.46 ms 1653.86 ms −0.40%
TP2+SP2 12.701 s 12.876 s +1.38% 1451.47 ms 1485.73 ms +2.36%
路径 旧版稳定阶段 chunk 间隔 新版稳定阶段 chunk 间隔 变化
单卡 stepwise 2741.08 ms 2742.31 ms +0.045%
TP2 1843.49 ms 1856.06 ms +0.68%
SP2 1893.71 ms 1778.75 ms −6.07%
TP2+SP2 1291.51 ms 1293.20 ms +0.13%

每个 stepwise 配置的 8 条请求、80 个媒体 chunk 全部一致;普通路径 8 条请求的最终帧数组哈希一致。以上配置的三个指标均未出现超过 3% 工程容差的退化;这不等于严格零退化,也不是统计等效性证明。

5.2 同一新版的跨请求共享开关对照

使用关闭 → 开启 → 开启 → 关闭的顺序。每次启动运行 4 条请求,排除首条冷请求,每种状态累计 6 条暖请求。计时期间关闭中间 tensor 记录。

拓扑 关闭共享首 chunk 开启共享首 chunk 变化 关闭共享完整耗时 开启共享完整耗时 变化
TP2 1715.79 ms 1818.95 ms +6.01% 17.6805 s 17.8081 s +0.72%
SP2 1658.33 ms 1850.59 ms +11.59% 17.0885 s 17.3071 s +1.28%
TP2+SP2 1466.49 ms 1777.22 ms +21.19% 12.7165 s 13.0272 s +2.44%

每种拓扑的 16 条请求、160 个媒体 chunk 输出均一致。

拓扑 关闭共享进程 RSS 开启共享进程 RSS 增加
TP2 8672.82 MiB 12275.77 MiB +3602.95 MiB
SP2 8767.11 MiB 12416.71 MiB +3649.59 MiB
TP2+SP2 15810.18 MiB 19412.26 MiB +3602.09 MiB

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. 支持范围与已知限制

  • 已实现的 BSND cross 契约覆盖合法 GQA、Q/KV/pair mask 和两种并行放置;原生 packed/varlen cross metadata 仍明确拒绝。H3 的 packed joint self-attention 沿用原有路径。
  • 跨请求首块共享要求 eager 执行及 clean-x0 K/V 提交,不支持的组合在初始化时拒绝。
  • 编译正确性声明限定为上述 TP2+SP2、关闭共享配置。
  • LingBot 当前 AR 请求路径不支持多个请求同时在 GPU 上批处理执行;其 max_num_seqs > 1 被拒绝。通用执行契约可按 BSND batch 行或显式 pair mask 隔离请求,但不等于 LingBot 已支持多请求 AR 批处理。
  • 跨请求首块存储当前为单条 CPU 快照。
  • 开启共享存在已测得的首 chunk 延迟和 CPU 内存开销。
  • 有限次数的取消和资源测试不能证明长期或细小泄漏不存在。
  • 本 PR 不重写视频分页管理,不新增显存分配器,也不声明所有模型及后端已接入。DreamZero 的选择性窗口重置缓存规则、Sana 接入及 LTX2 原生模型集成仍是后续工作;编译模式与首块跨请求共享组合仍不支持。

AI assistance: 使用 Codex 辅助实现、测试执行、代码整理及 PR 正文撰写。生成视频与首帧已在上文标注。

…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>
@chatgpt-codex-connector

Copy link
Copy Markdown

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

@vllm-omni-review-bot

Copy link
Copy Markdown

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.

@tsinghua-code

Copy link
Copy Markdown
Contributor Author

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):

  • Checked that the diff contains only the 10 production files covering the generic cross-attention contract, K/V cache lifecycle management, and LingBot reference integration.
  • Checked explicit cross semantics, request-scoped conditioning K/V isolation, guarded publication/invalidation, and first-block snapshot restoration into request-owned video pages.
  • On the current submission, 39 targeted CPU tests and all applicable local pre-commit checks passed, including mypy.
  • Earlier validation covered 82 CPU tests, 2 CUDA lifecycle tests, and real-checkpoint comparisons of denoising outputs, latents, and video frames. Each complete request used 10 chunks. These results are tied to the tested revision documented in the PR; the full model matrix has not been rerun on the final submission.
  • Cross-request sharing remains disabled by default. Its measured first-chunk latency regressions and the unsupported compilation-plus-sharing combination are explicitly documented.
    AI assistance was used for implementation, validation, and preparing this summary.

Signed-off-by: tsinghua-code <zr-liu25@mails.tsinghua.edu.cn>
Signed-off-by: tsinghua-code <zr-liu25@mails.tsinghua.edu.cn>

tzhouam commented Oct 13, 2026

Copy link
Copy Markdown
Collaborator

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:

Model Requirement exposed by its implementation
Helios Already separates project_kv() from attention execution and accepts precomputed K/V. A good second reference integration; its tensor-backed cache should not require an AR paging manager.
Wan2.2 / DreamZero Text and image attention use separate softmaxes, then combine outputs. The shared abstraction must preserve this composition. DreamZero also retains text while refreshing image conditioning across window resets.
Sana-Video / Sana-WM Need consistent KV-padding semantics and shared backend dispatch instead of model-specific mask fallback.
Stable Audio Uses GQA but currently expands K/V heads in model code. Native head geometry and backend-owned GQA adaptation belong in the common contract.
LTX2 Requires both sharded-Q/replicated-KV and replicated-Q/sharded-KV execution. Its audio/video cross-attention K/V changes with the denoising state and cannot be treated as static conditioning.
MiniMax H3 Uses packed joint text/audio/video self-attention. It should reuse the attention infrastructure while preserving joint computation and document isolation.

Based on these cases, I would refine the design as follows:

  1. Keep projected Q/K/V execution as the shared boundary. Models retain projections, normalization, RoPE, query selection, and multi-condition composition. In particular, summing independent text/image attention outputs must not become attention over concatenated K/V.

  2. Separate computation semantics, parallel placement, and cache policy. computation="cross" should neither force skip_sequence_parallel=True nor imply that K/V is reusable. Reuse the existing parallel-strategy interface for communication.

  3. Integrate the contract with backend capabilities. Extend the existing execution context and metadata to describe head geometry, mask semantics, and packed boundaries. Backend capability resolution should replace backend-name allowlists. Perform mask-content validation during input preparation to avoid per-layer host synchronization.

  4. Separate cache validity from storage. Preserve generation-checked publication, invalidation, stream readiness, and cleanup, while allowing both tensor-backed storage and the existing AR manager. Cache identity needs component-level condition revisions, projection ownership/revisions, logical CFG branches, and adapter state. Preserve Helios’s source-lifetime protection.

  5. Keep first-block sharing as an explicit snapshot policy. It remains part of this PR, with its own ownership, dependency matching, and performance evidence. The attention executor should not depend on that policy.

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.

@hsliuustc0106 hsliuustc0106 added core related to core module: cache, scheduler, engine, worker, modelrunner diffusion codes related to diffusion models high priority high priority issue, needs to be done asap labels Oct 14, 2026
Signed-off-by: tsinghua-code <zr-liu25@mails.tsinghua.edu.cn>
@tsinghua-code

Copy link
Copy Markdown
Contributor Author

@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 6f55ecc2f. The shared boundary remains execution of already projected Q/K/V; model code still owns projections, normalization, RoPE, query selection, and composition of independent conditioning branches.

  1. Computation, placement, and cache policy are separate. computation="cross" is independent of the configuration role. It neither forces skip_sequence_parallel nor marks K/V reusable. The common parallel-strategy interface now covers local Q with replicated K/V and replicated Q with sequence-sharded K/V; the latter also supports GQA and masks. LTX2 itself is not migrated.
  2. Backend capability checks now describe the actual request. The execution context carries Q/KV head geometry, mask kind, and cross placement. Q validity, KV validity, and Q–KV pair restrictions are distinct. The selected backend either handles them, uses the documented SDPA fallback, or rejects an unsupported explicit configuration. Mask content is checked when metadata is prepared, rather than by synchronizing on each layer call. Stable Audio now passes native GQA heads to the common layer, and Wan2.2 declares cross computation explicitly. Text/image branches still execute separate softmaxes before the model combines their outputs.
  3. Cache validity is separate from storage. ConditioningKVScope owns identity, guarded complete publication, invalidation, CUDA readiness, and idempotent cleanup. It works with LingBot's existing AR storage and new request-owned tensor storage used by Helios. Identity includes condition/projection revisions, projection owner, logical branch, adapter revision, and storage representation. Helios also checks the exact source tensor plus its mutation version or a snapshot when a version is unavailable; this retains the protection against the [Bugfix][Helios] Guard cross-attn KV cache against address reuse #8064 address-reuse failure. Text conditioning K/V stays request-local.
  4. First-block sharing remains an independent, opt-in snapshot policy. A hit requires the full first-block dependency set; a borrowed CPU snapshot is restored into a new request's private GPU video pages. The attention executor has no dependency on the sharing store.

Validation summary (old and new use the same checkpoint, prompt, seed, and execution path within each row):

Validation Result
Integrated LingBot ordinary / stepwise Each generated 10 chunks. Old/new: 40/40 denoise and 40/40 latent records identical in each path; ordinary 117 decoded frames identical; stepwise 10/10 video chunks and media bytes identical.
Integrated Helios ordinary / stepwise Each generated 10 chunks. Old/new: 66/66 denoise, 66/66 scheduler latent, 10/10 decode latent, and 10/10 frame blocks bitwise identical in each path. The Helios baseline already has explicit cross execution, so this comparison specifically validates its cache integration.
Focused contract and lifecycle 99 current-contract checks passed in the combined worktree, including invalid GQA geometry, automatic SDPA fallback, and the accepted cross-SP placement. Earlier checks on the unchanged shared attention files covered SP2 (6 cases) and TP2+SP2 (8 cases), all with zero numerical error, plus eight exact old/new Helios/Wan/Stable Audio module comparisons on SDPA and FlashAttention.
Helios lifecycle and performance Real-checkpoint cancellation/reuse loop: 25 rounds, no retained request tensors or allocated-memory growth in the measured steady rounds; real engine cancellation and failure recovery also passed. In an old/new/old restart comparison, median 10-chunk time was 40.0020 s old vs 39.7288 s new (−0.68%, small sample).

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.

@tzhouam tzhouam added the ready label to trigger buildkite CI label Oct 16, 2026
@tzhouam
tzhouam self-requested a review October 16, 2026 03:22

@tzhouam tzhouam left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@tzhouam
tzhouam enabled auto-merge (squash) October 16, 2026 07:24
@tzhouam tzhouam added ready label to trigger buildkite CI and removed ready label to trigger buildkite CI labels Oct 16, 2026
Signed-off-by: tsinghua-code <zr-liu25@mails.tsinghua.edu.cn>
auto-merge was automatically disabled October 16, 2026 10:35

Head branch was pushed to by a user without write access

@tzhouam tzhouam added ready label to trigger buildkite CI and removed ready label to trigger buildkite CI labels Oct 16, 2026
@vllm-omni-review-bot

Copy link
Copy Markdown

Omni ReviewBot: CI is red on this head

@tsinghua-code required checks failed on 1afdb3c2bb37:

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.

@tzhouam
tzhouam enabled auto-merge (squash) October 17, 2026 05:50
@tzhouam
tzhouam merged commit 69d298a into vllm-project:main Oct 17, 2026
11 checks passed
princepride added a commit to Levius-Fubuki/vllm-omni that referenced this pull request Oct 18, 2026
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>
MrlixiangWE added a commit to MrlixiangWE/vllm-omni that referenced this pull request Oct 21, 2026
…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>
shihongzhi added a commit to shihongzhi/vllm-omni that referenced this pull request Oct 21, 2026
…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>
@tsinghua-code
tsinghua-code deleted the codex/track-a-8659-production-only branch October 29, 2026 07:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core related to core module: cache, scheduler, engine, worker, modelrunner diffusion codes related to diffusion models high priority high priority issue, needs to be done asap ready label to trigger buildkite CI

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants