From 8287d4eb73b9c90e97ec751b9d1bf3af38a1ada4 Mon Sep 17 00:00:00 2001 From: ActivePeter <1020401660@qq.com> Date: Mon, 27 Jul 2026 13:20:15 +0800 Subject: [PATCH 1/2] perf(lingbot): decouple condition prefetch from control ingress - prefetch two bounded condition chunks before control arrival - refill conditions after denoise to overlap with decode - anchor scheduler latency at control acceptance - enforce ordered control submission and add regression tests - document the final pipeline and timing comparison Verified: - 1045 unit tests passed, 13 skipped - 4-GPU chunk mean improved from 1.8010s to 1.4476s --- docs/zh/design_lingbot_condition_prefetch.md | 266 ++++++++++++++++++ docs/zh/stream_scheduler.md | 3 + .../streaming_pipeline_orchestrator.py | 11 +- .../pipelines/lingbot_world_fast/service.py | 6 +- .../pipelines/lingbot_world_fast/streaming.py | 114 ++++++-- .../test_streaming_pipeline_orchestrator.py | 21 ++ .../test_service_action_loop.py | 23 ++ .../lingbot_world_fast/test_streaming.py | 78 +++++ 8 files changed, 502 insertions(+), 20 deletions(-) create mode 100644 docs/zh/design_lingbot_condition_prefetch.md diff --git a/docs/zh/design_lingbot_condition_prefetch.md b/docs/zh/design_lingbot_condition_prefetch.md new file mode 100644 index 0000000..5a058c9 --- /dev/null +++ b/docs/zh/design_lingbot_condition_prefetch.md @@ -0,0 +1,266 @@ +# LingBot-World-Fast Condition 预取流水设计 + +> 状态:当前实现设计说明。本文描述内部调度策略,不新增用户配置、环境变量或服务协议。 + +## 1. 最终状态:完整运行时序 + +```mermaid +sequenceDiagram + autonumber + participant C as 控制客户端 + participant S as LingBot Service + participant P as Pipeline Runtime + participant O as Streaming Orchestrator + participant E as VAE Encode Actor + participant D as DiT Actor + participant V as VAE Decode Actor + + C->>S: 建立流式 session + S->>P: _create_initialized_session(config) + P->>P: encode prompt、准备图像、计算 latent/KV 几何 + P->>E: initialize encode cache(image, cache_handle) + E-->>P: encode cache ready + P->>V: initialize decode cache(cache_handle) + V-->>P: decode cache ready + P->>D: initialize denoise/KV/noise cache(cache_handle) + D-->>P: denoise cache ready + P-->>S: initialized runtime + + S->>O: create_session(runtime, final_sequence_id) + O-)E: enqueue condition[0] + O-)E: enqueue condition[1] + O-->>S: streaming session ready + + loop 每个 chunk i + par 与 control 无关的 condition 路径 + E-->>O: condition[i] ready + and 真实交互控制路径 + C->>S: control[i] + S->>P: resolve + validate control[i] + P-->>S: control tensor[i] + S->>O: try_submit control[i] + Note over S,O: control[i] 被接受时记录 latency anchor + end + + O->>O: 按 session/sequence join condition[i] + control[i] + + alt World-KV 已命中 latent[i] + O->>D: advance noise cursor + D-->>O: cached latent[i] + else 正常生成 + O->>D: denoise(condition[i], control[i], session caches) + D-->>O: latent[i] + end + + par Post-denoise condition refill + opt i+2 未越界且预取窗口有容量 + O-)E: enqueue condition[i+2] + end + and 当前 chunk decode + O->>V: decode latent[i] + V-->>O: frames[i] + end + + O->>O: 按 sequence ID 有序提交输出与 scheduler metrics + O-->>S: frames[i] + S-->>C: chunk[i] + applied controls + target facts + end + + C->>S: stop / disconnect / session complete + S->>O: close_session(drain) + O->>V: release decode cache + V-->>O: released + O->>D: release denoise/KV/noise cache + D-->>O: released + O->>E: release encode cache + E-->>O: released + O-->>S: session state released +``` + +图中的 `par` 表示两条路径之间没有数据依赖,不承诺它们在物理设备上同时执行。condition 可能早于 control 完成, +也可能在 control 到达后才完成;DiT 始终等待同一 sequence ID 的两项输入都就绪。 + +最终结构中有三条必须分开的路径: + +- `condition`:只依赖 session 图像、VAE causal state 和 chunk 位置,可以在 control 之前计算; +- `control`:来自真实交互输入,必须按 chunk 严格有序,不能预测或跨 chunk 复用; +- `latent/frames`:只有同一 sequence ID 的 condition 与 control join 后才能产生。 + +VAE Encode、DiT 和 VAE Decode 分别由唯一 actor 管理。它们即使放在同一张 GPU 上也可以重叠;调度器不会根据 +device placement 隐式创建 resource group。session 结束时,cache 通过 owning actor 按 +`decode → denoise → encode` 的逆拓扑顺序释放。 + +## 2. 最终状态的准入与状态边界 + +每个 session 维护两个单调递增 cursor: + +- `next_condition_index`:下一条尚未提交的 condition 序号; +- `next_control_index`:下一条允许接收的 control 序号。 + +预取窗口满足: + +```text +0 <= next_condition_index - next_control_index <= 2 +``` + +control 准入必须同时满足: + +1. `chunk_index == next_control_index`; +2. 对应 condition 已提交,或当前仍可补交该 condition; +3. control edge 仍有容量。 + +只有 scheduler 接受 control 后,control cursor 才递增。重复、跳号和乱序 control 直接失败。 +`_ensure_next_condition_locked()` 只在预取被 backpressure 暂时耗尽时补齐当前 control 所需的一个 condition; +正常窗口补充由当前 chunk denoise 完成后的 refill 触发。 + +预取深度 `2` 是内部常量,不是用户配置。condition、control、latent 和 output edge 也都具有显式的 per-session +容量,因此长 session 不会形成 duration-sized tensor 列表。Condition actor 自身仍然串行执行;两级预取表示 +最多允许一个任务执行、另一个任务或结果停留在有界路径中,不表示同一个 VAE worker 同时运行两个 encode。 + +## 3. 本次 Diff:优化前后时序对比 + +### 3.1 优化前 + +```mermaid +sequenceDiagram + autonumber + participant C as 控制客户端 + participant S as LingBot Service + participant O as Streaming Orchestrator + participant E as VAE Encode Actor + participant D as DiT Actor + participant V as VAE Decode Actor + + C->>S: first control[0] + Note over S: 收到首个 control 后才初始化 + S->>S: create runtime + session + + loop 每个 chunk i + S->>O: atomic push(encode_request[i], control[i]) + O->>E: encode condition[i] + E-->>O: condition[i] + O->>D: condition[i] + control[i] + D-->>O: latent[i] + O->>V: decode latent[i] + V-->>O: frames[i] + O-->>S: frames[i] + S-->>C: chunk[i] + end + + Note over C,V: control 等待、condition encode、denoise、decode 更容易形成串行链 +``` + +优化前,`encode_request` 与 `control` 必须在一次原子 ingress 中提交。condition 明明与当前 control 无关,却必须 +等待 control 到达;首个 control 之后还要承担 runtime/session 初始化,直接拉长第一条交互关键路径。 + +### 3.2 优化后 + +```mermaid +sequenceDiagram + autonumber + participant C as 控制客户端 + participant S as LingBot Service + participant O as Streaming Orchestrator + participant E as VAE Encode Actor + participant D as DiT Actor + participant V as VAE Decode Actor + + S->>S: create runtime + session + S->>O: create_session(runtime) + O-)E: prefetch condition[0] + O-)E: prefetch condition[1] + + par Condition 可以提前完成 + E-->>O: condition[0] + and 等待真实控制 + C->>S: control[0] + S->>O: submit control[0] + Note over S,O: 从 control acceptance 开始计时 + end + + O->>D: join(condition[0], control[0]) + D-->>O: latent[0] + + par 后置补充未来 condition + O-)E: prefetch condition[2] + and 解码当前 latent + O->>V: decode latent[0] + V-->>O: frames[0] + end + + O-->>S: ordered frames[0] + S-->>C: chunk[0] +``` + +### 3.3 变化分析 + +| 维度 | 优化前 | 优化后 | +|---|---|---| +| 首个 control | 先等待 control,再初始化 runtime/session | 先初始化并预取,再等待 control | +| Condition ingress | 与 control 原子提交 | 与 control 解耦,最多领先 2 个 chunk | +| DiT 启动条件 | encode 在 control 后开始,随后 join | condition 可提前就绪,control 到达后即可 join | +| 后续 refill | 下一 control 到达后才触发 encode | 当前 denoise 完成后补充窗口 | +| 主要重叠 | 重叠机会少,容易形成串行链 | 未来 VAE encode 主要与当前 VAE decode 重叠 | +| Control 顺序 | 依赖原子 ingress 隐式维持 | `next_control_index` 显式拒绝重复、跳号和乱序 | +| 延迟起点 | 第一条任意 ingress | `latency_anchor_artifact="control"` | +| 内存边界 | edge 容量有界 | 保持有界,并增加固定深度 2 的 condition lookahead | + +Post-denoise refill 是调度偏好,不是“VAE encode 永不与 DiT 重叠”的硬互斥保证。如果 control 已到达而匹配的 +condition 尚未提交,活性保护仍可补交该 condition,以免 session 停滞。 + +## 4. 指标语义 + +Condition 预取可能早于 control 数秒。如果仍使用第一条 ingress 的时间作为起点,scheduler 会把用户尚未发送 +control 的等待时间错误计入 control-to-output latency。 + +`StreamingPipelineSpec.latency_anchor_artifact="control"` 规定: + +- condition-only ingress 不写入 `ingress_accepted_at`; +- control 被接受时才记录该 sequence 的 ingress 时间; +- 对应输出发出后,scheduler 才计算 control-to-output latency。 + +该字段默认是 `None`,其他 pipeline 保持“第一条任意 ingress”为起点的旧行为。它只修正 scheduler 指标语义, +不会合并 target chunk compute、客户端 delivery latency 或 AIPerf warmup/聚合职责。 + +## 5. 收益与代价 + +以下数据来自固定 4×H100、`chunk_size=3`、SageAttention SM90 的累计 checkpoint;它们不是严格的单开关 A/B: + +| Checkpoint | Chunk mean | 相对上一阶段 | 相对初始基线 | +|---|---:|---:|---:| +| 4 卡 SageAttention 基线 | 1.800984s | - | - | +| 首轮 condition/cache 版本 | 1.695177s | -5.9% | -5.9% | +| Condition 与 control 解耦预取 | 1.579448s | -6.8% | -12.3% | +| Post-denoise refill | 1.449961s | -8.2% | -19.5% | + +按相邻累计 checkpoint 看,post-denoise refill 是当时最大的单项下降;按正确性依赖看,应把 runtime 前置、 +condition/control 解耦、两级预取、后置 refill、control 顺序和延迟锚点作为一个整体。 + +主要代价:连接建立后即使用户一直不发送 control,也会提前占用 runtime、三类 session cache 和最多两个 +condition slot。断连、停止、stage failure 或预取失败时,仍由原有 session cleanup 释放这些状态。 + +本设计不改变扩散步数、模型权重、dtype、attention 实现、VAE 数值路径或输出顺序,属于无损调度优化。 + +## 6. 实现映射 + +| 设计职责 | 实现位置 | +|---|---| +| 首个 control 前创建 runtime/session | `telefuser/pipelines/lingbot_world_fast/service.py::_run_actor_worker_loop` | +| Session cursor、两级预取和 control 准入 | `telefuser/pipelines/lingbot_world_fast/streaming.py` | +| Post-denoise refill | `LingBotWorldFastStreamingRuntime::_denoise` | +| Control latency anchor | `telefuser/orchestrator/streaming_pipeline_orchestrator.py` | +| 时序、顺序与指标回归 | `tests/unit/pipelines/lingbot_world_fast/`、`tests/unit/orchestrator/` | + +## 7. 验证要求 + +功能测试至少覆盖: + +- session 创建后、control 到达前已提交两个 condition; +- denoise 完成后触发 refill,并补入下一 condition; +- 重复和乱序 control 被拒绝; +- condition-only ingress 不启动 control latency timer; +- 单 chunk、多 session、World-KV hit、stage failure 和 cleanup failure 行为不回退。 + +性能复测必须固定模型、分辨率、`chunk_size`、扩散 schedule、attention 实现、GPU 数量、placement 和 AIPerf +warmup 口径。正式结论至少比较多轮 mean、P90、P99、加权 compute FPS,并同时检查 GPU/VAE/DiT 时序,避免 +把单轮调度抖动误判为稳定收益。 diff --git a/docs/zh/stream_scheduler.md b/docs/zh/stream_scheduler.md index f0aaa23..155b694 100644 --- a/docs/zh/stream_scheduler.md +++ b/docs/zh/stream_scheduler.md @@ -102,3 +102,6 @@ LingBot 的 `vae_encode_config` 和 `vae_decode_config` 是两个独立且完整 - session state 必须隔离,并通过 owning actor 释放。 - 只为真实且明确的部署约束声明 resource group。 - 应验证 session 交错、backpressure、取消、actor failure 和 cleanup failure。 + +LingBot-World-Fast 的完整 condition/control 流水、最终时序和本次优化前后对比见 +[LingBot-World-Fast Condition 预取流水设计](design_lingbot_condition_prefetch.md)。 diff --git a/telefuser/orchestrator/streaming_pipeline_orchestrator.py b/telefuser/orchestrator/streaming_pipeline_orchestrator.py index 019f6db..e593454 100644 --- a/telefuser/orchestrator/streaming_pipeline_orchestrator.py +++ b/telefuser/orchestrator/streaming_pipeline_orchestrator.py @@ -141,10 +141,13 @@ class StreamingPipelineSpec: output_artifacts: frozenset[str] = frozenset() output_capacity_per_session: int = 1 resource_groups: tuple[StreamingResourceGroupSpec, ...] = () + latency_anchor_artifact: str | None = None def __post_init__(self) -> None: if self.output_capacity_per_session < 1: raise ValueError("output_capacity_per_session must be at least one") + if self.latency_anchor_artifact == "": + raise ValueError("latency_anchor_artifact must be non-empty when provided") @dataclass(frozen=True) @@ -561,7 +564,8 @@ def try_push_inputs(self, session_id: str, sequence_id: int, inputs: Mapping[str for edge in edges ): return False - runtime.ingress_accepted_at.setdefault(sequence_id, time.monotonic()) + if self.spec.latency_anchor_artifact is None or self.spec.latency_anchor_artifact in inputs: + runtime.ingress_accepted_at.setdefault(sequence_id, time.monotonic()) for artifact, value in inputs.items(): runtime.artifacts[(sequence_id, artifact)] = _ArtifactSlot( value, {edge.target_stage for edge in edges_by_artifact[artifact]} @@ -1033,6 +1037,11 @@ def _validate(self) -> None: for artifact in self.spec.output_artifacts: if artifact not in declared_producers: raise ValueError(f"Output artifact {artifact!r} has no producer") + if ( + self.spec.latency_anchor_artifact is not None + and self.spec.latency_anchor_artifact not in external_artifacts + ): + raise ValueError(f"Latency anchor artifact {self.spec.latency_anchor_artifact!r} is not an external input") stage_positions = {stage.stage_id: index for index, stage in enumerate(self.spec.stages)} roots = deque(stage.stage_id for stage in self.spec.stages if indegree[stage.stage_id] == 0) diff --git a/telefuser/pipelines/lingbot_world_fast/service.py b/telefuser/pipelines/lingbot_world_fast/service.py index 85d301c..8bb7dac 100644 --- a/telefuser/pipelines/lingbot_world_fast/service.py +++ b/telefuser/pipelines/lingbot_world_fast/service.py @@ -895,9 +895,6 @@ def _run_actor_worker_loop( emit_status: Callable[..., None], ) -> None: """Drive dynamic control ingress and ordered output through the shared actor graph.""" - first_item = self._next_realtime_control(state, control_context, control_builder, 0, emit_status, block=True) - if first_item is None: - return runtime_measurement = self._start_benchmark_measurement(state) try: runtime = self.pipeline._create_initialized_session(state.config, progress_callback=emit_status) @@ -921,6 +918,9 @@ def _run_actor_worker_loop( runtime=self._runtime_metadata(runtime), **({"measurement": {"name": "runtime_creation", **runtime_facts}} if runtime_facts is not None else {}), ) + first_item = self._next_realtime_control(state, control_context, control_builder, 0, emit_status, block=True) + if first_item is None: + return submitted = 0 controls_by_chunk: dict[int, list[str] | None] = {} diff --git a/telefuser/pipelines/lingbot_world_fast/streaming.py b/telefuser/pipelines/lingbot_world_fast/streaming.py index b7b5587..6b85860 100644 --- a/telefuser/pipelines/lingbot_world_fast/streaming.py +++ b/telefuser/pipelines/lingbot_world_fast/streaming.py @@ -33,6 +33,9 @@ from .pipeline import LingBotWorldFastPipeline +_CONDITION_PREFETCH_DEPTH = 2 + + @dataclass(frozen=True) class LingBotWorldFastStreamingSession: """Lightweight identity for one session in the shared streaming runtime.""" @@ -47,6 +50,8 @@ class _LingBotStreamingSessionEntry: runtime: LingBotWorldFastGenerationSession epoch: int progress_callback: Callable[..., None] | None + next_condition_index: int = 0 + next_control_index: int = 0 class LingBotWorldFastStreamingRuntime: @@ -90,6 +95,7 @@ def __init__(self, pipeline: LingBotWorldFastPipeline) -> None: ), output_artifacts=frozenset({"frames"}), output_capacity_per_session=2, + latency_anchor_artifact="control", ) self.orchestrator = StreamingPipelineOrchestrator(spec, actors) @@ -108,16 +114,28 @@ def create_session( if session_id in self._sessions: raise ValueError(f"LingBot streaming session {session_id!r} already exists") epoch = self.orchestrator.create_session(session_id, final_sequence_id=runtime.chunk_count - 1) - self._sessions[session_id] = _LingBotStreamingSessionEntry(runtime, epoch, progress_callback) - return LingBotWorldFastStreamingSession(session_id, epoch, runtime.cache_handle) + entry = _LingBotStreamingSessionEntry(runtime, epoch, progress_callback) + self._sessions[session_id] = entry + session = LingBotWorldFastStreamingSession(session_id, epoch, runtime.cache_handle) + try: + self._prefetch_conditions(session_id, entry) + except BaseException: + self.close_session(session) + raise + return session def can_submit_chunk(self, session: LingBotWorldFastStreamingSession) -> bool: - """Return whether the session can atomically admit another chunk.""" - self._require_session(session) + """Return whether the session can admit its next control chunk.""" + entry = self._require_session(session) try: - if self.orchestrator.status(session.session_id) != StreamingSessionStatus.RUNNING: - return False - return self.orchestrator.can_push_inputs(session.session_id, ("encode_request", "control")) + with self._lock: + if self.orchestrator.status(session.session_id) != StreamingSessionStatus.RUNNING: + return False + if entry.next_control_index >= entry.runtime.chunk_count: + return False + if not self._ensure_next_condition_locked(session.session_id, entry): + return False + return self.orchestrator.can_push_input(session.session_id, "control") except RuntimeError: if self.orchestrator.error(session.session_id) is not None: return False @@ -139,25 +157,88 @@ def try_submit_chunk( chunk_index: int, control: torch.Tensor, ) -> bool: - """Atomically submit one chunk to the shared actor graph.""" + """Submit the next control while condition encoding runs ahead independently.""" entry = self._require_session(session) if chunk_index < 0 or chunk_index >= entry.runtime.chunk_count: raise ValueError("chunk_index exceeds the LingBot session length") try: - return self.orchestrator.try_push_inputs( - session.session_id, - chunk_index, - { - "encode_request": None, - "control": control, - }, - ) + with self._lock: + if chunk_index != entry.next_control_index: + raise ValueError(f"Expected LingBot control chunk {entry.next_control_index}, got {chunk_index}") + if not self._ensure_next_condition_locked(session.session_id, entry): + return False + accepted = self.orchestrator.try_push_inputs( + session.session_id, + chunk_index, + {"control": control}, + ) + if accepted: + entry.next_control_index += 1 + return accepted except RuntimeError: error = self.orchestrator.error(session.session_id) if error is not None: raise RuntimeError("LingBot streaming scheduler failed") from error raise + def _prefetch_conditions( + self, + session_id: str, + entry: _LingBotStreamingSessionEntry, + ) -> None: + with self._lock: + self._prefetch_conditions_locked(session_id, entry) + + def _prefetch_conditions_locked( + self, + session_id: str, + entry: _LingBotStreamingSessionEntry, + ) -> None: + """Keep at most two condition chunks ahead of accepted controls.""" + while ( + entry.next_condition_index < entry.runtime.chunk_count + and entry.next_condition_index - entry.next_control_index < _CONDITION_PREFETCH_DEPTH + ): + if not self.orchestrator.try_push_inputs( + session_id, + entry.next_condition_index, + {"encode_request": None}, + ): + return + entry.next_condition_index += 1 + + def _ensure_next_condition_locked( + self, + session_id: str, + entry: _LingBotStreamingSessionEntry, + ) -> bool: + """Ensure the next control has a matching condition without filling the lookahead window.""" + if entry.next_condition_index > entry.next_control_index: + return True + if entry.next_condition_index >= entry.runtime.chunk_count: + return False + if not self.orchestrator.try_push_inputs( + session_id, + entry.next_condition_index, + {"encode_request": None}, + ): + return False + entry.next_condition_index += 1 + return True + + def _refill_conditions_after_denoise( + self, + session_id: str, + entry: _LingBotStreamingSessionEntry, + ) -> None: + """Overlap future condition encoding with decode instead of current-chunk DiT work.""" + with self._lock: + if self._sessions.get(session_id) is not entry: + return + if self.orchestrator.status(session_id) != StreamingSessionStatus.RUNNING: + return + self._prefetch_conditions_locked(session_id, entry) + def poll_frames(self, session: LingBotWorldFastStreamingSession) -> list[tuple[int, list[Image.Image]]]: """Return decoded frame batches in chunk order.""" self._require_session(session) @@ -388,6 +469,7 @@ def _denoise(self, invocation: StreamingStageInvocation) -> dict[str, object]: runtime.world_kv_binding.on_chunk_finalized(runtime, index, latent) except Exception as exc: logger.warning(f"world_kv on_chunk_finalized failed at chunk {index}: {exc}") + self._refill_conditions_after_denoise(invocation.key.session_id, entry) return {"latent": latent} def _decode_inputs(self, invocation: StreamingStageInvocation) -> tuple[tuple[object, ...], dict[str, object]]: diff --git a/tests/unit/orchestrator/test_streaming_pipeline_orchestrator.py b/tests/unit/orchestrator/test_streaming_pipeline_orchestrator.py index 4163cb1..9d0d233 100644 --- a/tests/unit/orchestrator/test_streaming_pipeline_orchestrator.py +++ b/tests/unit/orchestrator/test_streaming_pipeline_orchestrator.py @@ -45,6 +45,7 @@ def _wait_for_outputs( def _orchestrator( encode_calls: list[tuple[int, bool, bool]], denoise_calls: list[tuple[int, bool, bool]], + latency_anchor_artifact: str | None = None, ) -> StreamingPipelineOrchestrator: def encode(invocation): encode_calls.append((invocation.key.sequence_id, invocation.is_first, invocation.is_last)) @@ -65,6 +66,7 @@ def denoise(invocation): StreamingEdgeSpec("control", "denoise"), ), output_artifacts=frozenset({"frames"}), + latency_anchor_artifact=latency_anchor_artifact, ) return StreamingPipelineOrchestrator( spec, @@ -127,6 +129,25 @@ def test_session_metrics_report_end_to_end_percentiles_after_output_polling() -> orchestrator.close() +def test_session_metrics_use_the_configured_ingress_latency_anchor() -> None: + orchestrator = _orchestrator([], [], latency_anchor_artifact="control") + try: + orchestrator.create_session("session", final_sequence_id=0) + orchestrator.push_input("session", 0, "image", "prefetched") + assert orchestrator.wait_until_idle("session") + assert orchestrator.session_metrics("session").ingress_accepted_at == () + + control_submitted_after = time.monotonic() + orchestrator.push_input("session", 0, "control", "left") + assert _wait_for_outputs(orchestrator, "session", 1) + metrics = orchestrator.session_metrics("session") + assert len(metrics.ingress_accepted_at) == 1 + assert metrics.ingress_accepted_at[0][1] >= control_submitted_after + assert metrics.control_to_output_latency.count == 1 + finally: + orchestrator.close() + + def test_bounded_condition_edge_applies_backpressure_until_downstream_consumes() -> None: encode_calls: list[tuple[int, bool, bool]] = [] denoise_calls: list[tuple[int, bool, bool]] = [] diff --git a/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py b/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py index bd1365f..b1e9659 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py @@ -157,6 +157,29 @@ def initialize(*_args: object, **_kwargs: object) -> LingBotWorldFastGenerationS assert state.generation_session is runtime +def test_actor_worker_prefetches_conditions_before_waiting_for_first_control() -> None: + pipeline = MagicMock() + runtime = LingBotWorldFastGenerationSession(config=_state().config, latent_f=1, chunk_size=1, cache_handle=7) + pipeline._create_initialized_session.return_value = runtime + streaming_runtime = MagicMock() + streaming_runtime.create_session.return_value = SimpleNamespace(session_id="actor-session") + pipeline._get_streaming_runtime.return_value = streaming_runtime + service = LingBotWorldFastService(pipeline) + state = _state() + + def stop_after_prefetch(*_args: object, **_kwargs: object) -> None: + assert pipeline._create_initialized_session.called + assert streaming_runtime.create_session.called + state.active = False + return None + + with patch.object(service, "_next_realtime_control", side_effect=stop_after_prefetch): + service._run_actor_worker_loop(state, MagicMock(), MagicMock(), MagicMock()) + + streaming_runtime.create_session.assert_called_once() + streaming_runtime.try_submit_chunk.assert_not_called() + + def test_actor_worker_prefetches_directional_chunks_within_ingress_capacity() -> None: pipeline = MagicMock() runtime = LingBotWorldFastGenerationSession( diff --git a/tests/unit/pipelines/lingbot_world_fast/test_streaming.py b/tests/unit/pipelines/lingbot_world_fast/test_streaming.py index 2e81a34..79497ac 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_streaming.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_streaming.py @@ -52,12 +52,20 @@ class _Denoise: def __init__(self, release_order: list[str]): self.release_order = release_order self.advance_calls: list[int] = [] + self.calls: list[dict[str, object]] = [] self.release_calls: list[int] = [] self.fail_denoise = False + self.started: threading.Event | None = None + self.release: threading.Event | None = None def denoise_and_update_cache(self, **kwargs): if self.fail_denoise: raise RuntimeError("injected denoise failure") + self.calls.append(kwargs) + if self.started is not None: + self.started.set() + if self.release is not None: + self.release.wait(timeout=1) return kwargs["condition_chunk"] def advance_noise(self, cache_handle: int): @@ -161,6 +169,76 @@ def test_streaming_session_routes_one_chunk_through_three_stages() -> None: assert runtime.cache_handle is None +def test_streaming_session_prefetches_two_conditions_ahead_of_controls() -> None: + pipeline = _Pipeline() + runtime = LingBotWorldFastGenerationSession( + config=LingBotWorldFastSessionConfig(prompt="test", image=Image.new("RGB", (8, 8))), + prompt_emb=torch.tensor([0.0]), + latent_h=1, + latent_w=1, + latent_f=4, + height=8, + width=8, + frame_tokens=1, + chunk_size=1, + max_attention_size=1, + cache_handle=15, + ) + streaming_runtime = LingBotWorldFastStreamingRuntime(pipeline) + session = streaming_runtime.create_session(runtime) + denoise_started = threading.Event() + release_denoise = threading.Event() + pipeline.denoise_stage.started = denoise_started + pipeline.denoise_stage.release = release_denoise + try: + assert streaming_runtime.wait_until_idle(session) + assert [call["chunk_index"] for call in pipeline.vae_encode_worker.calls] == [0, 1] + assert pipeline.denoise_stage.calls == [] + + streaming_runtime.submit_chunk(session, 0, torch.tensor([4.0])) + assert denoise_started.wait(timeout=1) + assert [call["chunk_index"] for call in pipeline.vae_encode_worker.calls] == [0, 1] + release_denoise.set() + assert streaming_runtime.wait_until_idle(session) + assert [call["chunk_index"] for call in pipeline.vae_encode_worker.calls] == [0, 1, 2] + assert len(pipeline.denoise_stage.calls) == 1 + assert streaming_runtime.poll_frames(session)[0][0] == 0 + metrics = streaming_runtime.session_metrics(session) + assert [index for index, _ in metrics.ingress_accepted_at] == [0] + finally: + release_denoise.set() + streaming_runtime.close_session(session) + streaming_runtime.close() + + +def test_streaming_session_rejects_out_of_order_or_duplicate_controls() -> None: + pipeline = _Pipeline() + runtime = LingBotWorldFastGenerationSession( + config=LingBotWorldFastSessionConfig(prompt="test", image=Image.new("RGB", (8, 8))), + prompt_emb=torch.tensor([0.0]), + latent_h=1, + latent_w=1, + latent_f=2, + height=8, + width=8, + frame_tokens=1, + chunk_size=1, + max_attention_size=1, + cache_handle=16, + ) + streaming_runtime = LingBotWorldFastStreamingRuntime(pipeline) + session = streaming_runtime.create_session(runtime) + try: + with pytest.raises(ValueError, match="Expected LingBot control chunk 0, got 1"): + streaming_runtime.try_submit_chunk(session, 1, torch.tensor([4.0])) + assert streaming_runtime.try_submit_chunk(session, 0, torch.tensor([4.0])) + with pytest.raises(ValueError, match="Expected LingBot control chunk 1, got 0"): + streaming_runtime.try_submit_chunk(session, 0, torch.tensor([4.0])) + finally: + streaming_runtime.close_session(session) + streaming_runtime.close() + + def test_streaming_runtime_shares_one_actor_graph_across_sessions() -> None: pipeline = _Pipeline() streaming_runtime = LingBotWorldFastStreamingRuntime(pipeline) From 33967f951cdeaf70de42e55e9e7a0cfb1fdce85e Mon Sep 17 00:00:00 2001 From: ActivePeter <1020401660@qq.com> Date: Tue, 28 Jul 2026 10:33:52 +0800 Subject: [PATCH 2/2] fix(lingbot): harden condition prefetch lifecycle Changes: - keep capacity polling side-effect free and atomically admit missing conditions with controls - run denoising in inference mode and close example pipelines on every post-load exit path - add regression coverage and document the condition prefetch design and performance Verification: - git diff --cached --check - pytest -q tests/unit/test_run_examples.py tests/unit/pipelines/lingbot_world_fast/test_streaming.py (blocked during collection: torch is not installed) --- docs/zh/blog_lingbot_condition_prefetch.md | 459 ++++++++++++++++++ docs/zh/design_lingbot_condition_prefetch.md | 4 +- examples/run_examples.py | 189 ++++---- .../pipelines/lingbot_world_fast/streaming.py | 41 +- .../lingbot_world_fast/test_streaming.py | 50 ++ tests/unit/test_run_examples.py | 88 ++++ 6 files changed, 714 insertions(+), 117 deletions(-) create mode 100644 docs/zh/blog_lingbot_condition_prefetch.md create mode 100644 tests/unit/test_run_examples.py diff --git a/docs/zh/blog_lingbot_condition_prefetch.md b/docs/zh/blog_lingbot_condition_prefetch.md new file mode 100644 index 0000000..7acf83c --- /dev/null +++ b/docs/zh/blog_lingbot_condition_prefetch.md @@ -0,0 +1,459 @@ +# TeleFuser世界模型推理优化记录:简单的overlap优化了蛮多 + +如果一个世界模型已经能持续生成视频,它就算“实时交互”了吗?不一定。 + +持续生成只回答了“画面能不能不断输出”。交互还要求用户发出一个控制后,下一段画面尽快体现这个控制。方向已经 +改变,画面却过了一两段才响应,帧率再稳定,体感仍然会像高延迟的云游戏。 + +下面先梳理交互式世界模型从建立 session 到输出视频 chunk 的完整流程,再从数据依赖中定位延迟来源和优化切入点。 + +## 世界模型一次交互推理做了什么 + +LingBot-World-Fast 不是收到一条 control 就独立生成一段视频的无状态服务。它是一台持续运行、带有时间状态的 +生成器:前一个 chunk 的 VAE 状态、干净 latent 和 KV cache 都会影响下一个 chunk。 + +一次完整 session 可以分为三个阶段。 + +### 1. Session 初始化 + +用户提供 prompt、初始图像、seed 和生成参数后,系统先建立后续所有 chunk 共享的运行状态: + +- 文本编码器把 prompt 转成 prompt embedding; +- 图像被缩放、归一化,并确定 latent 的时空尺寸; +- 根据 seed 初始化 noise generator; +- VAE Encode 和 VAE Decode 分别建立自己的 causal cache; +- DiT 分配 self-attention KV cache、cross-attention cache 和调度器状态。 + +这些工作通常只在 session 开始时做一次,但它们决定了第一条 control 最早何时能进入模型。 + +### 2. 按 Chunk 循环生成 + +每个 chunk 都要接收用户控制、生成 latent,再解码成画面。DiT 内部还要执行多个扩散 timestep,并在最后把干净 +状态写回 KV cache,供未来 chunk 使用。 + +### 3. Session 结束与状态释放 + +用户停止、断连、生成完成或某个 stage 失败后,系统必须排空或取消任务,并释放 VAE、DiT 和调度器持有的 +session 状态。否则一次失败就可能留下几十 GiB 显存或长期存活的 worker。 + +把这三个阶段连起来,整体时序如下: + +```mermaid +sequenceDiagram + autonumber + participant U as 用户/控制端 + participant S as LingBot Service + participant R as Session Runtime + participant E as VAE Encode + participant D as DiT + KV Cache + participant V as VAE Decode + + U->>S: 建立 session(prompt, image, seed, config) + S->>R: 编码 prompt、准备图像、计算 latent 几何 + R->>E: 初始化 causal encode cache + E-->>R: encode cache ready + R->>V: 初始化 causal decode cache + V-->>R: decode cache ready + R->>D: 初始化 noise、scheduler、self/cross KV cache + D-->>R: denoise state ready + R-->>S: runtime ready + + loop 每个 chunk i + U->>S: action / camera control[i] + S->>S: 解析并校验 control[i] + R->>E: 编码 condition[i] + E-->>R: condition[i] + R->>D: prompt + condition[i] + control[i] + noise[i] + loop 多个 diffusion timestep + D->>D: DiT forward、预测 x0、必要时重新加噪 + end + D->>D: t=0 clean-cache commit + Note over D: 把干净历史写入 KV,供后续 chunk 使用 + D-->>V: clean latent[i] + V->>V: causal decode + V-->>S: frames[i] + S-->>U: 视频 chunk[i] + end + + U->>S: stop / disconnect + S->>V: 释放 decode state + S->>D: 释放 denoise/KV/noise state + S->>E: 释放 encode state + S-->>U: session closed +``` + +图中画的是逻辑处理顺序,不表示所有方框都必须在物理设备上串行执行。恰恰相反,调度优化的机会就藏在这些箭头 +里:有些箭头代表真正的数据依赖,有些只是旧实现选择了“等上一步一起提交”。 + +为了先把计算职责讲清楚,上图按“初始化 → chunk 循环”的逻辑阶段展开。旧调度实际上把初始化推迟到了第一条 +control 到达之后,这也是后文要处理的第一个问题。 + +## 理解关键路径:哪些工作真的必须等待 Control + +把复杂的模型细节压缩成数据依赖,会得到下面这张图: + +```mermaid +flowchart LR + P[Prompt] --> PE[Prompt Embedding] + I[初始图像 + VAE 状态] --> CE[Condition Encode] + C[用户 Control] --> CP[Control 解析与校验] + N[Noise + Scheduler + 历史 KV] --> J[DiT Denoise] + PE --> J + CE --> J + CP --> J + J --> K[Clean-cache Commit] + K --> VD[VAE Decode] + VD --> O[视频 Chunk] +``` + +先把图中的几个名词拆开: + +| 名词 | 在本文中指什么 | +| --- | --- | +| `chunk` | 一次流水线处理的一小段连续视频;同一个 session 中的 chunk 按 `0, 1, 2, ...` 排序。 | +| `prompt embedding` | 文本编码器对用户 prompt 生成的特征,作为 DiT 的文本语义条件。session 内 prompt 不变,因此只需计算一次。 | +| `condition[i]` | **数据。**传给 DiT 的第 `i` 个画面条件张量,由 mask 和 VAE latent 拼接而成。它是一次 VAE Encode 的产物,不是会自行推进的状态。 | +| `encode_request[i]` | **调度信号。**请求 VAE Encode Actor 为第 `i` 个 chunk 计算 `condition[i]`;它本身不包含 condition tensor。 | +| Condition 预取 | **调度动作。**在 `control[i]` 到达前提前提交 `encode_request[i]`,并允许编码结果在有界路径中等待。 | +| VAE causal state | **Session 状态。**VAE 编码前面 chunk 后保留的时序缓存,由 VAE Encode 按 chunk 顺序读取和更新。 | +| `control` | 当前 chunk 的用户控制,经相机位姿、内参以及可选 action 转换得到的模型输入。它描述接下来希望如何移动或操作,只有用户输入后才能确定。 | +| `noise` | 当前 chunk 开始扩散时使用的随机 latent。它由 session 的 seed 和随机数生成器状态依次产生,不能脱离 session 随意重排。 | +| 历史 `KV` | DiT self-attention 为已生成的干净 chunk 保存的 key/value cache。它让当前 chunk 能关注此前画面,并在每个 chunk 生成后继续更新。 | +| `sequence ID` | 调度器为 session 内每个 chunk 使用的序号,也就是上面的 `0, 1, 2, ...`。它是配对 `condition[i]`、`control[i]` 和输出的调度标识,不是模型生成的内容。 | + +三者的关系是:`encode_request[i]` 触发一次 VAE Encode;VAE Encode 读取并更新该 session 的 causal state,产出 +`condition[i]`。causal state 使后续编码延续同一段视频的时间上下文,而不是把每个 chunk 当作互不相关的独立 +视频。DiT 则是实际执行扩散去噪的 Transformer,只有它同时拿到当前 chunk 所需的输入后才能开始计算。 + +```mermaid +flowchart LR + R["encode_request[i]
调度信号"] --> E["VAE Encode
编码动作"] + S[("VAE causal state
Session 状态")] -->|读取| E + E -->|更新| S + E --> C["condition[i]
输出数据"] + + C --> J["按 sequence ID i
等待输入汇合"] + U["control[i]
用户输入数据"] --> J + J --> D["DiT Denoise"] + P["prompt embedding"] --> D + G["noise + 历史 KV"] --> D + D --> L["latent[i]"] +``` + +因此,上图中的依赖关系可以具体理解为: + +- `prompt embedding` 是 session 级输入,不需要每个 chunk 重算; +- VAE Encode 必须按 chunk 顺序推进 causal state,但提交 `encode_request[i]` 不需要等待用户的 `control[i]`; +- `control` 来自用户,不能预测;即使相邻 chunk 的控制数值碰巧相同,也必须分别对应各自的用户意图; +- `noise` 和历史 `KV` 属于 session 的生成状态,必须由该 session 自己持有; +- 对 chunk `i`,调度器可以提前提交 `encode_request[i]`,由 VAE Encode Actor 产出 `condition[i]`,同时独立等待 + `control[i]`。两项数据只需在 DiT 开始前以相同的 sequence ID `i` 汇合,彼此没有直接依赖。 + +最后一条就是本次优化的切入点。 + +## 世界模型的“快”其实有三种 + +在讨论优化前,还要区分三类经常被混在一起的时间: + +| 指标 | 回答的问题 | +| --- | --- | +| Session startup | 从建立连接到 runtime 可用要多久? | +| Control-to-output latency | control 被接受后,多久能看到体现它的画面? | +| Chunk period | 相邻两个视频 chunk 的输出间隔是否稳定? | + +外层 FPS 只能描述最终产出速度,不能单独回答控制反馈是否及时。一个系统可能平均 FPS 尚可,但第一条 control +等待了很久;也可能首段很快,后续 chunk period 却追不上播放消耗。 + +本次改动主要缩短前两项中的非必要等待。 + +## 找到切入点:控制来了,系统才开始备菜 + +三个指标中,这里的首要优化目标是 **Control-to-output latency**:把不依赖当前 control 的工作移出“control 被接受 +到对应画面输出”这条关键路径。提前创建 runtime 还会同时改善 **Session startup**,避免第一条 control 承担初始化 +成本;**Chunk period** 不是本节的直接优化目标,而是优化后仍需保持稳定的约束。 + +### 先直观理解 Condition 预取 + +Condition 预取做的事情很单纯:系统在等待用户提交 `control[i]` 时,先提交 `encode_request[i]`。Orchestrator +随后异步调度 VAE Encode Actor;Actor 读取该 session 的 causal state,计算 `condition[i]`,并让结果暂存在有界 +路径中。系统没有猜测用户控制,也没有提前调用 DiT 生成视频。 + +没有预取时,每轮都由 control 触发: + +```mermaid +flowchart LR + C["control[n] 到达"] --> A["原子提交
control[n] + encode_request[n]"] + A --> E["VAE Encode Actor"] + E --> J["condition[n] 与 control[n]
按 sequence ID n 汇合"] + A --> J + J --> D["DiT Denoise[n]"] + D --> V["VAE Decode
latent[n]"] + V --> O["输出 chunk[n]"] + O --> W["n ← n+1
等待下一条 control"] + W -.-> C +``` + +提前预取后,Denoise 完成会循环补充未来请求: + +```mermaid +flowchart LR + S["streaming session 创建
提交 encode_request[0]、[1]"] --> Q["有界 encode_request 路径"] + Q --> E["VAE Encode Actor
执行 encode_request[n]"] + E --> J["condition[n] 已就绪
control[n] 到达即汇合"] + C["control[n] 到达"] -->|最快路径| J + J --> D["DiT Denoise[n]"] + D --> V["VAE Decode
latent[n]"] + V --> O["输出 chunk[n]"] + + D --> R["完成后触发 refill
提交 encode_request[n+2]"] + R --> Q + + C -.-> F["请求缺失时
原子补交 encode_request[n]"] + F --> Q + + classDef fastPath fill:#dcfce7,stroke:#16a34a,stroke-width:3px,color:#14532d + class C,J,D fastPath +``` + +因此,预取没有减少 VAE Encode 的计算量,而是让这段计算尽量发生在用户尚未发出 control 的等待期内。如果 +`condition[i]` 已经就绪,`control[i]` 到达后可以直接进入 DiT;如果它还在编码,DiT 仍会等待两项输入全部就绪。 +本文最多允许两个 condition 编码任务或结果领先 control,不会一次性提交整个 session 的编码任务。 + +第一张图中,每一轮都必须先等 `control[n]`,再把它与 `encode_request[n]` 一起提交;输出后进入下一轮并继续等待。 +第二张图中的 `n` 表示当前正在处理的 chunk。session 创建时先提交 `0` 和 `1`;进入稳定循环后,`Denoise[n]` 完成 +会触发补窗,正常情况下提交 `encode_request[n+2]`。该请求被异步投递到 VAE Encode Actor,因此它的编码可能与 +当前 `latent[n]` 的 VAE Decode 重叠。如果请求因 backpressure 等原因没有提前进入流水线,`control[n]` 到达时 +还会触发图中的虚线兜底路径。绿色节点标出预取命中时的最快路径:condition 已在汇合点等待,control 被接受后 +可以直接进入 DiT,不再经过一次 VAE Encode。 + +第一张图就是旧流程:`encode_request` 和 `control` 被作为一组输入原子提交。这样做很直观,也容易保证顺序, +但它把两条本来独立的路径绑在了一起:即使 condition 完全不依赖当前 control,也必须等用户操作后才能开始编码。 + +这就像餐厅必须等客人说“少放辣”,才开始洗菜。调味确实依赖客人的选择,洗菜却不依赖。把两件事绑在一起, +客人看到的等待时间自然会变长。 + +第一条 control 更吃亏:旧流程在收到它之后才创建 runtime 和 session,首个交互延迟还要额外承担 prompt 编码、 +图像准备和三类 cache 初始化成本。 + +## 我们的改动:让两条路在 DiT 门口汇合 + +核心思路现在就很自然了: + +> `condition[i]` 不依赖 `control[i]`,所以应提前提交它的 `encode_request[i]`;编码结果只需在 DiT 入口与 +> `control[i]` 按 chunk 序号汇合。 + +具体改动有三项: + +1. 建立 session 后立即准备 runtime,不再等第一条 control。 +2. 创建 streaming session 时,立即提交最前面的两个 `encode_request`(不足两个时按实际 chunk 数量提交)。 +3. 当前 chunk 的 denoise 计算完成后,从 `next_condition_index` 开始提交未来编码请求,直到补满深度为 `2` 的窗口。 + +### 谁发起预取,谁执行编码 + +这里没有一个专门轮询队列的“预取 Actor”。预取由 `LingBotWorldFastStreamingRuntime` 持有的游标和事件触发: + +- session 创建时,Runtime 主动提交最前面的两个 `encode_request`; +- chunk `i` 的 denoise 计算完成后,Denoise Actor 调用 Runtime 的补窗逻辑;如果仍有容量,Runtime 从 + `next_condition_index` 开始提交请求,直到窗口补满;正常逐 chunk 推进时,第一条通常是 `encode_request[i+2]`; +- 如果 `control[i]` 到达时对应请求还未提交,`try_submit_chunk()` 会将 `encode_request[i]` 与 `control[i]` 作为一次 + 原子 ingress 兜底提交。 + +Orchestrator 收到 `encode_request` 后检查 stage 和队列容量,再异步投递给 VAE Encode Actor。VAE Encode Actor 只 +负责顺序执行编码并产出 `condition[i]`,不负责决定何时预取。补窗发生在 denoise 计算结束之后,随后当前 +`latent[i]` 进入 VAE Decode,因此未来 condition 的编码更有机会与当前 chunk 的解码重叠。 + +把 session 初始化、condition 预取、逐 chunk 生成和状态释放串起来,完整时序如下: + +```mermaid +sequenceDiagram + autonumber + participant C as 控制客户端 + participant S as LingBot Service + participant P as LingBot Runtime + participant O as Streaming Orchestrator + participant E as VAE Encode Actor + participant D as DiT Actor + participant V as VAE Decode Actor + + C->>S: 建立流式 session + S->>P: 创建并初始化 runtime + P->>P: 编码 prompt、准备图像、计算 latent/KV 几何 + P->>E: 初始化 encode cache(image, cache_handle) + E-->>P: encode cache ready + P->>V: 初始化 decode cache(cache_handle) + V-->>P: decode cache ready + P->>D: 初始化 denoise/KV/noise cache(cache_handle) + D-->>P: denoise cache ready + P-->>S: initialized runtime + + S->>P: 创建 streaming session + P->>O: create_session(runtime, final_sequence_id) + P->>O: 提交 encode_request[0]、encode_request[1](若存在) + O-)E: 按序异步执行两个 encode_request + P-->>S: streaming session ready + + loop 每个 chunk i + par 与 control 无关的 condition 路径 + E-->>O: condition[i] ready + and 真实交互控制路径 + C->>S: control[i] + S->>P: 解析并校验 control[i] + P-->>S: control tensor[i] + S->>O: 提交 control[i] + end + + O->>O: 按 session/sequence join condition[i] + control[i] + + O->>D: denoise(condition[i], control[i], session caches) + loop 多个 diffusion timestep + D->>D: DiT forward、预测 x0、必要时重新加噪 + end + D->>D: t=0 clean-cache commit + D->>P: denoise[i] 完成,触发 refill + opt i+2 未越界且预取窗口有容量 + P->>O: 提交 encode_request[i+2] + O-)E: 异步执行 encode_request[i+2] + end + D-->>O: latent[i] + + O->>V: decode latent[i] + Note over E,V: 未来 condition 的编码可能与当前 chunk 的 decode 重叠 + V-->>O: frames[i] + + O->>O: 按 sequence ID 有序提交输出与 scheduler metrics + O-->>S: frames[i] + S-->>C: chunk[i] + applied controls + target facts + end + + C->>S: stop / disconnect / session complete + S->>O: close_session(drain) + O->>V: 释放 decode cache + V-->>O: released + O->>D: 释放 denoise/KV/noise cache + D-->>O: released + O->>E: 释放 encode cache + E-->>O: released + O-->>S: session state released +``` + +图中的 `par` 表示两条路径没有直接的数据依赖,不承诺它们一定在物理设备上并行执行。如果 condition 已经准备好, +control 到达后就可以直接进入 DiT,而不用先等待一次 VAE encode。 + +这次优化没有修改模型权重、扩散步数、dtype、attention 实现或 VAE 数值路径。变化发生在任务何时进入流水线, +而不是张量如何计算。它带来的不是“模型算得更快”,而是减少本来可以并行、却被排成串行的等待。 + +## 为什么只预取两个,而不是越多越好 + +预取不是免费的。每个提前提交的 `encode_request` 及其产出的 condition tensor 都会占用队列位置、tensor 引用 +和 session cache。如果把整个长 session 的编码请求一次性全部提交,短期延迟可能好看,内存却会随视频时长不断 +增长。 + +TeleFuser 使用固定深度 `2` 的窗口,并为每条 tensor 路径设置 per-session 容量。一个 condition 正在编码时, +最多再允许一个任务或结果停留在有界路径中。 + +每个 session 维护两个简单的游标: + +- `next_condition_index`:下一条尚未提交 `encode_request` 的 condition 序号; +- `next_control_index`:下一条允许接收的 control。 + +它们始终满足: + +```text +0 <= next_condition_index - next_control_index <= 2 +``` + +这既限制了预取规模,也让 control 的顺序变得明确。重复、跳号或乱序 control 会直接被拒绝,而不是悄悄污染 +有状态 cache。 + +## 为什么在 Denoise 之后补充窗口 + +只在 session 开始时提交两个 `encode_request` 还不够。流水线向前推进后,窗口必须持续补充。 + +这里的“补充”不是修改某个 condition 状态,而是提交下一条尚未进入流水线的 `encode_request`。chunk `i` 的 +Denoise Actor 完成模型计算和 clean-cache commit 后,会调用 Runtime 的补窗逻辑。Runtime 检查两个游标和队列 +容量;满足条件时,它从 `next_condition_index` 开始向 Orchestrator 提交请求,直到窗口补满。正常逐 chunk 推进时, +第一条通常是 `encode_request[i+2]`,随后由 Orchestrator 异步投递给 VAE Encode Actor。 + +选择这个时机是为了避开当前 chunk 的 DiT 计算。DiT 通常是最重的阶段,如果 VAE Encode 恰好与它在同一 GPU 上 +争抢资源,理论上的并行可能反而增加抖动。请求在 denoise 结束后异步发出,而当前 `latent[i]` 随后进入 VAE +Decode,因此未来 condition 的编码更有机会与当前 chunk 的解码重叠。 + +这里说的是调度偏好,不是硬互斥。设备放置、CUDA 调度和 backpressure 都会影响实际重叠关系。如果当前 control +已经到达,而匹配的 `encode_request` 还没有进入流水线,系统仍会补交这条请求,避免为了“完美重叠”把 session +卡死。 + +这里还有一条容易被忽略的边界:`can_submit_chunk()` 这样的容量查询必须是纯读操作,不能因为 service 轮询一次 +就启动新的 VAE Encode。只有 `try_submit_chunk()` 真正携带 control 时,系统才允许兜底补交 `encode_request`, +而且缺失的请求与 control 会作为同一次原子 ingress 接受或拒绝,避免只提交其中一半。 + +## 实测结果:从 1.80 秒到 1.45 秒 + +固定环境为 4 张 H100、`chunk_size=3`、SageAttention SM90。以下是累计 checkpoint,不是把多个开关放在同一次 +进程内做的严格单变量 A/B: + +| Checkpoint | Chunk mean | 相对上一阶段 | 相对初始基线 | +| --- | ---: | ---: | ---: | +| 4 卡 SageAttention 基线 | 1.800984s | - | - | +| 首轮 condition/cache 版本 | 1.695177s | -5.9% | -5.9% | +| Condition 与 control 解耦预取 | 1.579448s | -6.8% | -12.3% | +| Post-denoise refill | 1.449961s | -8.2% | -19.5% | + +最新一次 AIPerf 验证得到: + +| 指标 | 结果 | +| --- | ---: | +| Mean | 1.4476s | +| P90 | 1.4903s | +| P99 | 1.5110s | +| 加权 compute FPS | 8.2898 | + +AIPerf Run ID: + +```text +0710537a-6302-41b2-a3b5-d1944ed7991f +``` + +从平均值看,累计下降约 `19.6%`。更值得关注的是 P90 和 P99:交互系统不仅要平均快,还要减少偶发的长等待。 + +不过,这还不能直接宣称“完全实时”。稳态 chunk 通常代表约 12 个输出帧;在 16 FPS 下,它对应约 `0.75s` 的 +播放时长。要让生成长期追上播放消耗,P95 chunk period 还应稳定低于 `0.75s`,并为编码和网络传输留出余量。 +目前的结果说明关键路径被明显缩短,但 DiT 本身仍是下一阶段的主要优化对象。 + +## 这次优化属于无损调度优化 + +这次修改没有减少任何模型计算,也没有使用近似结果。保持不变的内容包括: + +- 模型权重和 checkpoint; +- 扩散 schedule 与 timestep; +- dtype; +- attention 实现; +- VAE encode/decode 数值路径; +- control 与 condition 的同序号对应关系; +- 最终输出顺序。 + +改变的只有独立任务的提交时间和重叠方式。对同一个 seed、输入和 control,计算图中的数学操作没有被删减。 + +## 可以容忍的代价:提前预热资源 + +现在 session 建立后,即使用户暂时不发送 control,系统也会提前: + +- 创建 runtime; +- 初始化 VAE、DiT 和 KV/noise cache; +- 占用最多两个 condition slot。 + +这相当于在 session 级别做预热:先完成 runtime/cache 初始化和有界的 condition 准备,再等待真实 control。在线 +推理系统普遍使用预热把初始化成本移出请求关键路径;这里最多保留两个 condition slot,收益是更短的首次控制 +延迟,因此当前代价可以容忍。生产环境仍需要 session 空闲超时、断连清理和并发容量限制。 + +只有当系统进一步追求冷启动弹性,例如空闲时回收 runtime、按需分配 GPU 或 scale to zero,提前预热与资源弹性 +之间的矛盾才会变得突出。届时可以把轻量 session 建立、runtime 激活和空闲回收拆开评估,这会成为后续独立的 +优化点,而不是本次交互延迟优化需要解决的问题。 + +所有状态继续归属于单个 session,并由拥有对应 worker 的 actor 按 `decode → denoise → encode` 的逆拓扑顺序 +释放。预取失败、stage failure、主动停止和断连也走同一套 cleanup。 + +## 关于 TeleFuser + +TeleFuser 由中国电信人工智能研究院(TeleAI)世界模型团队和 Infra 团队共同研发,由中国电信首席科学家 +李学龙教授带领。TeleFuser 是一个面向实时世界模型的高性能推理框架,支持流式生成、持续状态管理和低延迟交互, +现已[开源](https://github.com/Tele-AI/TeleFuser)。 diff --git a/docs/zh/design_lingbot_condition_prefetch.md b/docs/zh/design_lingbot_condition_prefetch.md index 5a058c9..af7fea0 100644 --- a/docs/zh/design_lingbot_condition_prefetch.md +++ b/docs/zh/design_lingbot_condition_prefetch.md @@ -110,8 +110,8 @@ control 准入必须同时满足: 3. control edge 仍有容量。 只有 scheduler 接受 control 后,control cursor 才递增。重复、跳号和乱序 control 直接失败。 -`_ensure_next_condition_locked()` 只在预取被 backpressure 暂时耗尽时补齐当前 control 所需的一个 condition; -正常窗口补充由当前 chunk denoise 完成后的 refill 触发。 +如果预取因 backpressure 暂时耗尽,`try_submit_chunk()` 会把当前 control 与其缺失的 condition 作为同一次原子 +ingress 提交;正常窗口补充仍由当前 chunk denoise 完成后的 refill 触发。纯容量查询不会提交 condition。 预取深度 `2` 是内部常量,不是用户配置。condition、control、latent 和 output edge 也都具有显式的 per-session 容量,因此长 session 不会形成 duration-sized tensor 列表。Condition actor 自身仍然串行执行;两级预取表示 diff --git a/examples/run_examples.py b/examples/run_examples.py index 85aa064..eb4ef9c 100644 --- a/examples/run_examples.py +++ b/examples/run_examples.py @@ -1032,6 +1032,19 @@ def _emit_result(data: dict) -> None: print(f"{_RESULT_MARKER}{json.dumps(data)}", flush=True) +def _close_pipeline(pipeline: object | None) -> None: + """Release pipeline-owned workers without masking the regression result.""" + if pipeline is None: + return + close = getattr(pipeline, "close", None) + if not callable(close): + return + try: + close() + except Exception as exc: + print(f"Warning: failed to close pipeline: {exc}", file=sys.stderr, flush=True) + + def _run_single(pipeline_key: str, config_path: str | None, output_dir: str | None) -> None: """Subprocess entry point: load, run, save one pipeline. @@ -1102,106 +1115,104 @@ def _run_single(pipeline_key: str, config_path: str | None, output_dir: str | No ) sys.exit(1) - # Phase 2: Inference - output = None try: - if torch.cuda.is_available(): - torch.cuda.synchronize() - torch.cuda.reset_peak_memory_stats() + # Phase 2: Inference + output = None + try: + if torch.cuda.is_available(): + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + + output = _call_run(module, pipeline, runner_config) + + if torch.cuda.is_available(): + torch.cuda.synchronize() + gpu_mem_peak = torch.cuda.max_memory_allocated() / (1024 * 1024) + except Exception as e: + if torch.cuda.is_available(): + torch.cuda.synchronize() + gpu_mem_peak = torch.cuda.max_memory_allocated() / (1024 * 1024) + tb = traceback.format_exc() + category = "OOM_ERROR" if _is_oom(e) else "INFERENCE_ERROR" + _emit_result( + { + "status": "ERROR", + "error": f"{e}\n{tb}", + "error_category": category, + "elapsed": round(time.time() - start, 2), + "peak_gpu_memory_mb": round(gpu_mem_peak, 2), + } + ) + sys.exit(1) - output = _call_run(module, pipeline, runner_config) + # Phase 3: Validation & Save + warnings = _validate_output(output) + ppl_config = getattr(module, "PPL_CONFIG", {}) + num_steps = ppl_config.get("num_inference_steps") + if isinstance(num_steps, list): + num_steps = sum(num_steps) + output_fps = ppl_config.get("target_fps", 15) - if torch.cuda.is_available(): - torch.cuda.synchronize() - gpu_mem_peak = torch.cuda.max_memory_allocated() / (1024 * 1024) - except Exception as e: - if torch.cuda.is_available(): - torch.cuda.synchronize() - gpu_mem_peak = torch.cuda.max_memory_allocated() / (1024 * 1024) - tb = traceback.format_exc() - category = "OOM_ERROR" if _is_oom(e) else "INFERENCE_ERROR" - _emit_result( - { - "status": "ERROR", - "error": f"{e}\n{tb}", - "error_category": category, - "elapsed": round(time.time() - start, 2), - "peak_gpu_memory_mb": round(gpu_mem_peak, 2), - } - ) - del pipeline - gc.collect() - sys.exit(1) + # First save to temp location to get resolution + temp_dir = os.path.join(output_root, "temp", timestamp) + try: + temp_path, num_frames, resolution = _save_output(output, temp_dir, ppl_cfg.output_type, fps=output_fps) + except Exception as e: + tb = traceback.format_exc() + _emit_result( + { + "status": "ERROR", + "error": f"{e}\n{tb}", + "error_category": "OUTPUT_ERROR", + "elapsed": round(time.time() - start, 2), + "peak_gpu_memory_mb": round(gpu_mem_peak, 2), + } + ) + sys.exit(1) - # Phase 3: Validation & Save - warnings = _validate_output(output) - ppl_config = getattr(module, "PPL_CONFIG", {}) - num_steps = ppl_config.get("num_inference_steps") - if isinstance(num_steps, list): - num_steps = sum(num_steps) - output_fps = ppl_config.get("target_fps", 15) + # Move to final location with correct filename + final_filename = _generate_output_filename(ppl_cfg.script, ppl_cfg.gpu_count, resolution, ppl_cfg.output_type) + final_path = os.path.join(target_dir, final_filename) + + if temp_path and os.path.exists(temp_path): + shutil.move(temp_path, final_path) + # Clean up temp directory + try: + os.rmdir(temp_dir) + parent_temp = os.path.dirname(temp_dir) + if not os.listdir(parent_temp): + os.rmdir(parent_temp) + except OSError: + pass # Directory not empty or other error + + elapsed = time.time() - start + status = "PASS" + error_msg = "" + if warnings: + severe = [w for w in warnings if "NaN" in w or "Inf" in w or "is None" in w] + if severe: + status = "ERROR" + error_msg = "; ".join(warnings) - # First save to temp location to get resolution - temp_dir = os.path.join(output_root, "temp", timestamp) - try: - temp_path, num_frames, resolution = _save_output(output, temp_dir, ppl_cfg.output_type, fps=output_fps) - except Exception as e: - tb = traceback.format_exc() _emit_result( { - "status": "ERROR", - "error": f"{e}\n{tb}", - "error_category": "OUTPUT_ERROR", - "elapsed": round(time.time() - start, 2), + "status": status, + "output_path": final_path, + "error": error_msg, + "elapsed": round(elapsed, 2), "peak_gpu_memory_mb": round(gpu_mem_peak, 2), + "num_frames": num_frames, + "resolution": resolution, + "num_steps": num_steps, + "script": ppl_cfg.script, } ) - pipeline = None # noqa: F841 + finally: + _close_pipeline(pipeline) + pipeline = None gc.collect() - sys.exit(1) - - # Move to final location with correct filename - final_filename = _generate_output_filename(ppl_cfg.script, ppl_cfg.gpu_count, resolution, ppl_cfg.output_type) - final_path = os.path.join(target_dir, final_filename) - - if temp_path and os.path.exists(temp_path): - shutil.move(temp_path, final_path) - # Clean up temp directory - try: - os.rmdir(temp_dir) - parent_temp = os.path.dirname(temp_dir) - if not os.listdir(parent_temp): - os.rmdir(parent_temp) - except OSError: - pass # Directory not empty or other error - - elapsed = time.time() - start - status = "PASS" - error_msg = "" - if warnings: - severe = [w for w in warnings if "NaN" in w or "Inf" in w or "is None" in w] - if severe: - status = "ERROR" - error_msg = "; ".join(warnings) - - _emit_result( - { - "status": status, - "output_path": final_path, - "error": error_msg, - "elapsed": round(elapsed, 2), - "peak_gpu_memory_mb": round(gpu_mem_peak, 2), - "num_frames": num_frames, - "resolution": resolution, - "num_steps": num_steps, - "script": ppl_cfg.script, - } - ) - - del pipeline # noqa: F821 - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() + if torch.cuda.is_available(): + torch.cuda.empty_cache() # ============================================================================= diff --git a/telefuser/pipelines/lingbot_world_fast/streaming.py b/telefuser/pipelines/lingbot_world_fast/streaming.py index 6b85860..786c9ce 100644 --- a/telefuser/pipelines/lingbot_world_fast/streaming.py +++ b/telefuser/pipelines/lingbot_world_fast/streaming.py @@ -125,7 +125,7 @@ def create_session( return session def can_submit_chunk(self, session: LingBotWorldFastStreamingSession) -> bool: - """Return whether the session can admit its next control chunk.""" + """Return whether the next control can be admitted without mutating ingress.""" entry = self._require_session(session) try: with self._lock: @@ -133,9 +133,10 @@ def can_submit_chunk(self, session: LingBotWorldFastStreamingSession) -> bool: return False if entry.next_control_index >= entry.runtime.chunk_count: return False - if not self._ensure_next_condition_locked(session.session_id, entry): - return False - return self.orchestrator.can_push_input(session.session_id, "control") + artifacts = ["control"] + if entry.next_condition_index <= entry.next_control_index: + artifacts.append("encode_request") + return self.orchestrator.can_push_inputs(session.session_id, artifacts) except RuntimeError: if self.orchestrator.error(session.session_id) is not None: return False @@ -165,14 +166,20 @@ def try_submit_chunk( with self._lock: if chunk_index != entry.next_control_index: raise ValueError(f"Expected LingBot control chunk {entry.next_control_index}, got {chunk_index}") - if not self._ensure_next_condition_locked(session.session_id, entry): - return False + if entry.next_condition_index < entry.next_control_index: + raise RuntimeError("LingBot condition cursor fell behind the control cursor") + needs_condition = entry.next_condition_index == entry.next_control_index + inputs: dict[str, object] = {"control": control} + if needs_condition: + inputs["encode_request"] = None accepted = self.orchestrator.try_push_inputs( session.session_id, chunk_index, - {"control": control}, + inputs, ) if accepted: + if needs_condition: + entry.next_condition_index += 1 entry.next_control_index += 1 return accepted except RuntimeError: @@ -207,25 +214,6 @@ def _prefetch_conditions_locked( return entry.next_condition_index += 1 - def _ensure_next_condition_locked( - self, - session_id: str, - entry: _LingBotStreamingSessionEntry, - ) -> bool: - """Ensure the next control has a matching condition without filling the lookahead window.""" - if entry.next_condition_index > entry.next_control_index: - return True - if entry.next_condition_index >= entry.runtime.chunk_count: - return False - if not self.orchestrator.try_push_inputs( - session_id, - entry.next_condition_index, - {"encode_request": None}, - ): - return False - entry.next_condition_index += 1 - return True - def _refill_conditions_after_denoise( self, session_id: str, @@ -446,6 +434,7 @@ def _denoise_kwargs(self, invocation: StreamingStageInvocation) -> dict[str, obj "max_attention_size": runtime.max_attention_size, } + @torch.inference_mode() def _denoise(self, invocation: StreamingStageInvocation) -> dict[str, object]: entry = self._entry_for_invocation(invocation) runtime = entry.runtime diff --git a/tests/unit/pipelines/lingbot_world_fast/test_streaming.py b/tests/unit/pipelines/lingbot_world_fast/test_streaming.py index 79497ac..3eaa821 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_streaming.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_streaming.py @@ -53,12 +53,14 @@ def __init__(self, release_order: list[str]): self.release_order = release_order self.advance_calls: list[int] = [] self.calls: list[dict[str, object]] = [] + self.inference_modes: list[bool] = [] self.release_calls: list[int] = [] self.fail_denoise = False self.started: threading.Event | None = None self.release: threading.Event | None = None def denoise_and_update_cache(self, **kwargs): + self.inference_modes.append(torch.is_inference_mode_enabled()) if self.fail_denoise: raise RuntimeError("injected denoise failure") self.calls.append(kwargs) @@ -164,6 +166,7 @@ def test_streaming_session_routes_one_chunk_through_three_stages() -> None: assert pipeline.vae_encode_worker.release_calls == [9] assert pipeline.vae_decode_worker.release_calls == [9] assert pipeline.denoise_stage.release_calls == [9] + assert pipeline.denoise_stage.inference_modes == [True] assert pipeline.release_order == ["decode", "denoise", "encode"] assert pipeline.released == [] assert runtime.cache_handle is None @@ -211,6 +214,53 @@ def test_streaming_session_prefetches_two_conditions_ahead_of_controls() -> None streaming_runtime.close() +def test_capacity_poll_is_pure_and_control_fallback_is_atomic() -> None: + pipeline = _Pipeline() + runtime = LingBotWorldFastGenerationSession( + config=LingBotWorldFastSessionConfig(prompt="test", image=Image.new("RGB", (8, 8))), + prompt_emb=torch.tensor([0.0]), + latent_h=1, + latent_w=1, + latent_f=4, + height=8, + width=8, + frame_tokens=1, + chunk_size=1, + max_attention_size=1, + cache_handle=19, + ) + streaming_runtime = LingBotWorldFastStreamingRuntime(pipeline) + session = streaming_runtime.create_session(runtime) + denoise_started = threading.Event() + release_denoise = threading.Event() + pipeline.denoise_stage.started = denoise_started + pipeline.denoise_stage.release = release_denoise + try: + assert streaming_runtime.wait_until_idle(session) + assert [call["chunk_index"] for call in pipeline.vae_encode_worker.calls] == [0, 1] + + streaming_runtime.submit_chunk(session, 0, torch.tensor([4.0])) + assert denoise_started.wait(timeout=1) + streaming_runtime.submit_chunk(session, 1, torch.tensor([5.0])) + assert [call["chunk_index"] for call in pipeline.vae_encode_worker.calls] == [0, 1] + + assert streaming_runtime.can_submit_chunk(session) + assert [call["chunk_index"] for call in pipeline.vae_encode_worker.calls] == [0, 1] + + assert streaming_runtime.try_submit_chunk(session, 2, torch.tensor([6.0])) + deadline = time.monotonic() + 1 + while len(pipeline.vae_encode_worker.calls) < 3 and time.monotonic() < deadline: + time.sleep(0.01) + assert [call["chunk_index"] for call in pipeline.vae_encode_worker.calls] == [0, 1, 2] + + assert not streaming_runtime.try_submit_chunk(session, 3, torch.tensor([7.0])) + assert [call["chunk_index"] for call in pipeline.vae_encode_worker.calls] == [0, 1, 2] + finally: + release_denoise.set() + streaming_runtime.close_session(session) + streaming_runtime.close() + + def test_streaming_session_rejects_out_of_order_or_duplicate_controls() -> None: pipeline = _Pipeline() runtime = LingBotWorldFastGenerationSession( diff --git a/tests/unit/test_run_examples.py b/tests/unit/test_run_examples.py new file mode 100644 index 0000000..de0995a --- /dev/null +++ b/tests/unit/test_run_examples.py @@ -0,0 +1,88 @@ +from __future__ import annotations + +from pathlib import Path +from types import ModuleType + +import pytest + +import examples.run_examples as run_examples +from examples.run_examples import _close_pipeline + + +class _ClosablePipeline: + def __init__(self, *, fail: bool = False) -> None: + self.fail = fail + self.close_calls = 0 + + def close(self) -> None: + self.close_calls += 1 + if self.fail: + raise RuntimeError("close failed") + + +def test_close_pipeline_releases_owned_workers() -> None: + pipeline = _ClosablePipeline() + + _close_pipeline(pipeline) + + assert pipeline.close_calls == 1 + + +def test_close_pipeline_does_not_mask_regression_result(capsys: pytest.CaptureFixture[str]) -> None: + pipeline = _ClosablePipeline(fail=True) + + _close_pipeline(pipeline) + + assert pipeline.close_calls == 1 + assert "Warning: failed to close pipeline: close failed" in capsys.readouterr().err + + +@pytest.mark.parametrize("failure_site", ["validate", "filename", "move", "emit"]) +def test_run_single_closes_pipeline_for_all_post_load_failures( + failure_site: str, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + pipeline = _ClosablePipeline() + module = ModuleType("test_example") + module.PPL_CONFIG = {} + config = run_examples.Config( + output_root=str(tmp_path / "output-root"), + pipelines={"test": run_examples.PipelineConfig(script="test_example.py")}, + ) + output_dir = tmp_path / "results" + output_dir.mkdir() + temp_path = tmp_path / "temporary.mp4" + temp_path.write_bytes(b"video") + + def raise_lifecycle_error(*args: object, **kwargs: object) -> None: + del args, kwargs + raise RuntimeError(f"injected {failure_site} failure") + + monkeypatch.setattr(run_examples, "load_config", lambda _path: config) + monkeypatch.setattr(run_examples, "_import_example_module", lambda _path: module) + monkeypatch.setattr(run_examples, "_patch_ppl_config", lambda _module, _overrides: None) + monkeypatch.setattr(run_examples, "_call_get_pipeline", lambda _module, _config: pipeline) + monkeypatch.setattr(run_examples, "_call_run", lambda _module, _pipeline, _config: object()) + monkeypatch.setattr(run_examples.torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(run_examples, "_validate_output", lambda _output: []) + monkeypatch.setattr( + run_examples, + "_save_output", + lambda _output, _temp_dir, _output_type, fps: (str(temp_path), 1, "1x1"), + ) + monkeypatch.setattr(run_examples, "_generate_output_filename", lambda *_args: "result.mp4") + + if failure_site == "validate": + monkeypatch.setattr(run_examples, "_validate_output", raise_lifecycle_error) + elif failure_site == "filename": + monkeypatch.setattr(run_examples, "_generate_output_filename", raise_lifecycle_error) + elif failure_site == "move": + monkeypatch.setattr(run_examples.shutil, "move", raise_lifecycle_error) + else: + monkeypatch.setattr(run_examples, "_emit_result", raise_lifecycle_error) + + with pytest.raises(RuntimeError, match=f"injected {failure_site} failure"): + run_examples._run_single("test", None, str(output_dir)) + + assert pipeline.close_calls == 1