From ce83a77fc10b83c468f40a50a9536b7bfa631c3e Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Fri, 31 Jul 2026 09:04:17 +0000 Subject: [PATCH 01/11] test(bmk): support aiperf for sglang --- .gitignore | 1 + benchmarks/telefuser_aiperf/README.md | 53 ++- ...eam_sglang_lingbot_world_v2_4gpu_1min.json | 39 ++ .../run_sglang_lingbot_world_v2_4gpu.sh | 55 +++ .../scripts/run_stream_bench.sh | 4 +- ...g_lingbot_world_v2_benchmark_contract.yaml | 66 ++++ .../telefuser_aiperf/__init__.py | 8 +- .../telefuser_aiperf/payload.py | 23 ++ .../telefuser_aiperf/sglang_adapter.py | 362 ++++++++++++++++++ .../tests/test_sglang_adapter.py | 264 +++++++++++++ docs/en/benchmark_aiperf.md | 17 +- scripts/setup_aiperf.sh | 11 +- 12 files changed, 894 insertions(+), 9 deletions(-) create mode 100644 benchmarks/telefuser_aiperf/configs/stream_sglang_lingbot_world_v2_4gpu_1min.json create mode 100755 benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh create mode 100644 benchmarks/telefuser_aiperf/sglang_lingbot_world_v2_benchmark_contract.yaml create mode 100644 benchmarks/telefuser_aiperf/telefuser_aiperf/sglang_adapter.py create mode 100644 benchmarks/telefuser_aiperf/tests/test_sglang_adapter.py diff --git a/.gitignore b/.gitignore index 99e68f0..8604a05 100755 --- a/.gitignore +++ b/.gitignore @@ -162,3 +162,4 @@ telefuser/_version.py !examples/data/lingbot_world_fast/image.jpg !examples/data/lingbot_world_fast/poses.npy !examples/data/lingbot_world_fast/intrinsics.npy +.models \ No newline at end of file diff --git a/benchmarks/telefuser_aiperf/README.md b/benchmarks/telefuser_aiperf/README.md index 47e8130..9b3e222 100644 --- a/benchmarks/telefuser_aiperf/README.md +++ b/benchmarks/telefuser_aiperf/README.md @@ -162,9 +162,54 @@ session, joins its room, receives native video tracks, sends reliable controls o `tf.status` and bounded `tf.metrics` messages. It produces AIPerf's standard session results without requiring any LiveKit-specific changes in AIPerf. +## SGLang LingBot-World v2, Four GPUs + +This target uses SGLang's native MessagePack WebSocket endpoint and does not start LiveKit. It requires an SGLang +installation containing the `/v1/realtime_video/generate` endpoint and +`LingBotWorldCausalDMDPipeline`. In terminal 1, launch the four-GPU server: + +```bash +bash benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh +``` + +The launcher defaults to GPUs `0,1,2,3`, port `30000`, and model +`robbyant/lingbot-world-v2-14b-causal-fast-diffusers`. It uses an installed `sglang` command when available. +For the checked-in `work_dirs/sglang` source, set `SGLANG_PYTHON` to an environment installed with that source's +dependencies. Override the defaults when needed: + +```bash +SGLANG_BIN=/path/to/sglang \ +SGLANG_PYTHON=/path/to/sglang-env/bin/python \ +SGLANG_SOURCE_DIR=/path/to/sglang-source \ +SGLANG_MODEL_PATH=/path/to/lingbot-world-v2-14b-causal-fast-diffusers \ +SGLANG_CUDA_VISIBLE_DEVICES=4,5,6,7 \ + bash benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh +``` + +Wait until `http://127.0.0.1:30000/health` succeeds, then run AIPerf in terminal 2: + +```bash +bash benchmarks/telefuser_aiperf/scripts/run_stream_bench.sh \ + benchmarks/telefuser_aiperf/configs/stream_sglang_lingbot_world_v2_4gpu_1min.json +``` + +Artifacts are written below +`artifacts/telefuser_aiperf/stream_sglang_lingbot_v2_4gpu_1min/`. The adapter counts combined and split SGLang frame +batches, converts the shared keyboard trace to `camera_actions` state events, and maps scheduler, WebP encoding, +pacing, and WebSocket write timings into AIPerf's standard stream result. + +The launcher intentionally passes `--flow-shift 10`, matching the official LingBot-World v2 implementation and the +TeleFuser workload. The SGLang source default is `5`; use `SGLANG_FLOW_SHIFT=5` only to benchmark SGLang's default +behavior, and do not compare that run as a numerically equivalent model configuration. The workload also fixes four +DMD steps, 16 FPS, 60 chunks, a KV window of 18 latent frames plus a six-frame sink, WebP quality 95, and no output +pacing. SGLang permits one active realtime generation session, so the contract fixes concurrency to one. + ## Troubleshooting - `The pinned streaming-capable AIPerf or LiveKit is not installed`: rerun `bash scripts/setup_aiperf.sh`. +- `SGLang is not installed or SGLANG_BIN is invalid`: activate the SGLang environment or set `SGLANG_BIN` to its + executable. +- SGLang connection refused on port 30000: wait for model loading to finish and check `/health`. - Connection refused on port 8088: the TeleFuser process is still warming up or has exited; inspect terminal 2. - `0/1 succeeded` with zero received frames: confirm LiveKit is still running, restart the TeleFuser service, wait for one idle worker, and rerun the benchmark. @@ -204,7 +249,7 @@ tree. History failures do not silently fall back to an in-memory or file-only da configs/ Reproducible batch and streaming workloads data/ Prompt and control inputs scripts/ Batch and streaming launchers -telefuser_aiperf/ Source-loaded LiveKit adapter +telefuser_aiperf/ Source-loaded LiveKit and SGLang realtime adapters tests/ Adapter tests *_contract.yaml Target and transport capability contracts ``` @@ -219,9 +264,11 @@ AIPerf environment first, then run the checks from the repository root: PYTHONPATH=benchmarks/telefuser_aiperf \ .venv-aiperf/bin/python -m pytest \ - benchmarks/telefuser_aiperf/tests/test_livekit_adapter.py + benchmarks/telefuser_aiperf/tests/test_livekit_adapter.py \ + benchmarks/telefuser_aiperf/tests/test_sglang_adapter.py bash -n \ scripts/setup_aiperf.sh \ - benchmarks/telefuser_aiperf/scripts/run_stream_bench.sh + benchmarks/telefuser_aiperf/scripts/run_stream_bench.sh \ + benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh ``` diff --git a/benchmarks/telefuser_aiperf/configs/stream_sglang_lingbot_world_v2_4gpu_1min.json b/benchmarks/telefuser_aiperf/configs/stream_sglang_lingbot_world_v2_4gpu_1min.json new file mode 100644 index 0000000..683e679 --- /dev/null +++ b/benchmarks/telefuser_aiperf/configs/stream_sglang_lingbot_world_v2_4gpu_1min.json @@ -0,0 +1,39 @@ +{ + "contract": "benchmarks/telefuser_aiperf/sglang_lingbot_world_v2_benchmark_contract.yaml", + "server_url": "http://127.0.0.1:30000", + "mode": "bidirectional", + "task": "bidirectional", + "prompt": "walk forward through the scene", + "image_path": "examples/data/lingbot_world_fast/image.jpg", + "fps": 16, + "session_count": 1, + "warmup_sessions": 0, + "warmup_chunks": 1, + "session_duration_s": 240.0, + "stagger_s": 0.0, + "control_trace_path": "benchmarks/telefuser_aiperf/data/stream_lingbot_controls.json", + "request_extra": { + "size": "832x480", + "num_frames": 957, + "max_chunks": 60, + "num_inference_steps": 4, + "guidance_scale": 1.0, + "seed": 42, + "realtime_output_format": "webp", + "output_compression": 95, + "realtime_output_pacing": false, + "realtime_causal_sink_size": 6, + "realtime_causal_kv_cache_num_frames": 18 + }, + "transport": { + "connect_timeout_s": 60.0, + "message_timeout_s": 300.0, + "frame_timeout_s": 300.0, + "shutdown_timeout_s": 10.0, + "receive_audio": false + }, + "server_metrics": { + "enabled": false + }, + "artifacts_dir": "artifacts/telefuser_aiperf/stream_sglang_lingbot_v2_4gpu_1min" +} diff --git a/benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh b/benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh new file mode 100755 index 0000000..521938d --- /dev/null +++ b/benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" +cd "${ROOT_DIR}" + +SGLANG_BIN="${SGLANG_BIN:-}" +SGLANG_SOURCE_DIR="${SGLANG_SOURCE_DIR:-${ROOT_DIR}/work_dirs/sglang}" +SGLANG_PYTHON="${SGLANG_PYTHON:-}" +SGLANG_MODEL_PATH="${SGLANG_MODEL_PATH:-robbyant/lingbot-world-v2-14b-causal-fast-diffusers}" +SGLANG_HOST="${SGLANG_HOST:-127.0.0.1}" +SGLANG_PORT="${SGLANG_PORT:-30000}" +SGLANG_CUDA_VISIBLE_DEVICES="${SGLANG_CUDA_VISIBLE_DEVICES:-0,1,2,3}" +SGLANG_FLOW_SHIFT="${SGLANG_FLOW_SHIFT:-10}" + +if [[ -n "${SGLANG_BIN}" ]]; then + if ! command -v "${SGLANG_BIN}" >/dev/null 2>&1; then + echo "SGLANG_BIN is not executable: ${SGLANG_BIN}" >&2 + exit 1 + fi + sglang_command=("${SGLANG_BIN}") +elif command -v sglang >/dev/null 2>&1; then + sglang_command=("$(command -v sglang)") +elif [[ -n "${SGLANG_PYTHON}" && -x "${SGLANG_PYTHON}" \ + && -f "${SGLANG_SOURCE_DIR}/python/sglang/cli/main.py" ]]; then + export PYTHONPATH="${SGLANG_SOURCE_DIR}/python${PYTHONPATH:+:${PYTHONPATH}}" + sglang_command=("${SGLANG_PYTHON}" "-c" "from sglang.cli.main import main; main()") +else + echo "SGLang is unavailable. Set SGLANG_BIN, or set SGLANG_SOURCE_DIR and a compatible SGLANG_PYTHON." >&2 + exit 1 +fi + +IFS=',' read -r -a gpu_ids <<< "${SGLANG_CUDA_VISIBLE_DEVICES}" +if [[ ${#gpu_ids[@]} -ne 4 ]]; then + echo "SGLANG_CUDA_VISIBLE_DEVICES must contain exactly four comma-separated GPU IDs." >&2 + exit 2 +fi + +export CUDA_VISIBLE_DEVICES="${SGLANG_CUDA_VISIBLE_DEVICES}" +export SGLANG_LINGBOT_LAZY_VAE_ENCODE_BLACK_FRAMES="${SGLANG_LINGBOT_LAZY_VAE_ENCODE_BLACK_FRAMES:-60}" + +exec "${sglang_command[@]}" serve \ + --model-path "${SGLANG_MODEL_PATH}" \ + --pipeline-class-name LingBotWorldCausalDMDPipeline \ + --host "${SGLANG_HOST}" \ + --port "${SGLANG_PORT}" \ + --num-gpus 4 \ + --ulysses-degree 4 \ + --flow-shift "${SGLANG_FLOW_SHIFT}" \ + --dit-cpu-offload false \ + --text-encoder-cpu-offload false \ + --vae-config.use-parallel-decode true \ + --vae-config.parallel-decode-mode spatial \ + --enable-torch-compile false \ + "$@" diff --git a/benchmarks/telefuser_aiperf/scripts/run_stream_bench.sh b/benchmarks/telefuser_aiperf/scripts/run_stream_bench.sh index 4dda0af..222f22b 100755 --- a/benchmarks/telefuser_aiperf/scripts/run_stream_bench.sh +++ b/benchmarks/telefuser_aiperf/scripts/run_stream_bench.sh @@ -20,8 +20,8 @@ else fi if [[ -z "${AIPERF_PYTHON}" ]] || ! command -v "${AIPERF_PYTHON}" >/dev/null 2>&1 \ || ! PYTHONPATH="${ADAPTER_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" \ - "${AIPERF_PYTHON}" -c 'import livekit, telefuser_aiperf' >/dev/null 2>&1; then - echo "The pinned streaming-capable AIPerf or LiveKit is not installed. Run: bash scripts/setup_aiperf.sh" >&2 + "${AIPERF_PYTHON}" -c 'import livekit, msgspec, websockets, telefuser_aiperf' >/dev/null 2>&1; then + echo "The pinned streaming-capable AIPerf dependencies are not installed. Run: bash scripts/setup_aiperf.sh" >&2 exit 1 fi diff --git a/benchmarks/telefuser_aiperf/sglang_lingbot_world_v2_benchmark_contract.yaml b/benchmarks/telefuser_aiperf/sglang_lingbot_world_v2_benchmark_contract.yaml new file mode 100644 index 0000000..865f27a --- /dev/null +++ b/benchmarks/telefuser_aiperf/sglang_lingbot_world_v2_benchmark_contract.yaml @@ -0,0 +1,66 @@ +# Benchmark contract for SGLang's LingBot-World v2 realtime WebSocket target. +contract_version: v1 +name: sglang_lingbot_world_v2_realtime +mode: stream_world +implementation: sglang +model_family: lingbot_world_v2 +model: LingBot-World-v2-14B-Causal-Fast +supported_tasks: + - bidirectional +transport: websocket +adapter: sglang_realtime +transport_provider: sglang +endpoint: + health_path: /health + offer_path: /v1/realtime_video/generate +request_encoding: + content_type: application/msgpack + init_type: init + control_type: event + control_kind: camera_actions +result_delivery: + media: websocket_frame_batch + metrics: websocket_chunk_stats + event_log: events/{phase}_{logical_session_index}_{session_id}.jsonl +workload: + mode: bidirectional + task: bidirectional + fps: 16 + session_count: 1 + warmup_sessions: 0 + session_duration_s: 240.0 + control_trace: benchmarks/telefuser_aiperf/data/stream_lingbot_controls.json + request_extra: + size: 832x480 + num_frames: 957 + max_chunks: 60 + num_inference_steps: 4 + guidance_scale: 1.0 + seed: 42 + realtime_output_format: webp + output_compression: 95 + realtime_output_pacing: false + realtime_causal_sink_size: 6 + realtime_causal_kv_cache_num_frames: 18 +metrics: + - offer_rtt_ms + - connected_latency_ms + - first_frame_latency_ms + - first_metadata_latency_ms + - stream_fps + - session_runtime_s + - frames_received + - control_ack_latency_ms + - control_to_next_frame_latency_ms + - chunk_compute_seconds + - chunk_compute_fps + - chunk_encode_seconds + - chunk_output_write_seconds + - success_rate +limits: + active_sessions: 1 +artifacts: + config: benchmarks/telefuser_aiperf/configs/stream_sglang_lingbot_world_v2_4gpu_1min.json + control_trace: benchmarks/telefuser_aiperf/data/stream_lingbot_controls.json + server_runner: benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh + benchmark_runner: benchmarks/telefuser_aiperf/scripts/run_stream_bench.sh diff --git a/benchmarks/telefuser_aiperf/telefuser_aiperf/__init__.py b/benchmarks/telefuser_aiperf/telefuser_aiperf/__init__.py index c2fec9c..7834d47 100644 --- a/benchmarks/telefuser_aiperf/telefuser_aiperf/__init__.py +++ b/benchmarks/telefuser_aiperf/telefuser_aiperf/__init__.py @@ -5,6 +5,7 @@ from aiperf.streaming.adapters import register_stream_adapter from telefuser_aiperf.adapter import TeleFuserLiveKitAdapter +from telefuser_aiperf.sglang_adapter import SGLangRealtimeAdapter def register_adapters(*, replace: bool = False) -> None: @@ -15,6 +16,11 @@ def register_adapters(*, replace: bool = False) -> None: TeleFuserLiveKitAdapter, replace=replace, ) + register_stream_adapter( + "sglang_realtime", + SGLangRealtimeAdapter, + replace=replace, + ) -__all__ = ["TeleFuserLiveKitAdapter", "register_adapters"] +__all__ = ["SGLangRealtimeAdapter", "TeleFuserLiveKitAdapter", "register_adapters"] diff --git a/benchmarks/telefuser_aiperf/telefuser_aiperf/payload.py b/benchmarks/telefuser_aiperf/telefuser_aiperf/payload.py index 02cf65f..a9bbbef 100644 --- a/benchmarks/telefuser_aiperf/telefuser_aiperf/payload.py +++ b/benchmarks/telefuser_aiperf/telefuser_aiperf/payload.py @@ -2,6 +2,7 @@ from __future__ import annotations +from pathlib import Path from typing import Any from aiperf.streaming.models import StreamSessionPlan @@ -27,3 +28,25 @@ def build_telefuser_livekit_session_body( if plan.image_path: body["image_path"] = plan.image_path return body + + +def build_sglang_realtime_init(plan: StreamSessionPlan) -> dict[str, Any]: + """Build the MessagePack init request for SGLang realtime video.""" + + options = dict(plan.request_extra) + reserved = {"type", "prompt", "first_frame", "fps"} + conflicts = sorted(reserved.intersection(options)) + if conflicts: + raise ValueError("SGLang request_extra cannot override protocol fields: " + ", ".join(conflicts)) + if not plan.image_path: + raise ValueError("SGLang LingBot-World v2 requires image_path") + image_path = Path(plan.image_path) + if not image_path.is_file(): + raise FileNotFoundError(f"SGLang first frame does not exist: {image_path}") + return { + "type": "init", + "prompt": plan.prompt, + "first_frame": image_path.read_bytes(), + "fps": plan.fps, + **options, + } diff --git a/benchmarks/telefuser_aiperf/telefuser_aiperf/sglang_adapter.py b/benchmarks/telefuser_aiperf/telefuser_aiperf/sglang_adapter.py new file mode 100644 index 0000000..14d6e49 --- /dev/null +++ b/benchmarks/telefuser_aiperf/telefuser_aiperf/sglang_adapter.py @@ -0,0 +1,362 @@ +"""AIPerf adapter for SGLang realtime video WebSocket sessions.""" + +from __future__ import annotations + +import asyncio +import time +from collections.abc import Callable, Mapping +from pathlib import Path +from typing import Any +from urllib.parse import urlsplit, urlunsplit + +import msgspec +import websockets +from aiperf.common.redact import redact_string +from aiperf.streaming.adapters.common import StreamHttpClient, StreamTargetMetadataMixin +from aiperf.streaming.config import StreamProfileConfig +from aiperf.streaming.contracts import BenchmarkContract +from aiperf.streaming.events import StreamEventRecorder +from aiperf.streaming.models import ControlEventResult, SessionResult, StreamChunkMeasurement, StreamSessionPlan + +from telefuser_aiperf.payload import build_sglang_realtime_init + +_KEY_TO_ACTION = { + "ArrowUp": "w", + "ArrowDown": "s", + "ArrowLeft": "a", + "ArrowRight": "d", +} + + +def _websocket_url(server_url: str, endpoint: str) -> str: + parsed = urlsplit(server_url) + scheme = {"http": "ws", "https": "wss", "ws": "ws", "wss": "wss"}.get(parsed.scheme) + if scheme is None: + raise ValueError(f"Unsupported SGLang server URL scheme: {parsed.scheme}") + path = f"{parsed.path.rstrip('/')}/{endpoint.lstrip('/')}" + return urlunsplit((scheme, parsed.netloc, path, "", "")) + + +class _SGLangRealtimeSession: + def __init__(self, *, adapter: SGLangRealtimeAdapter, plan: StreamSessionPlan) -> None: + self.adapter = adapter + self.plan = plan + self.result = SessionResult( + logical_session_index=plan.logical_session_index, + phase=plan.phase, + mode=plan.mode, + planned_session_id=plan.planned_session_id, + session_id=plan.planned_session_id, + ) + self.events = StreamEventRecorder( + artifacts_dir=adapter.artifacts_dir, + phase=plan.phase, + logical_session_index=plan.logical_session_index, + planned_session_id=plan.planned_session_id, + print_events=adapter.config.print_events, + ) + self.started_at = 0.0 + self.active_started_at: float | None = None + self.first_frame_at: float | None = None + self.last_frame_at: float | None = None + self.first_frame_batch_size = 0 + self.first_metadata_at: float | None = None + self.websocket: Any = None + self.control_task: asyncio.Task[None] | None = None + self.control_sent_at: dict[int, float] = {} + self.control_by_event_id: dict[int, int] = {} + self.held_actions: set[str] = set() + self.pending_frame_header: Mapping[str, Any] | None = None + + async def run(self) -> SessionResult: + self.started_at = time.perf_counter() + self.events.record( + "session_start", + transport="websocket", + transport_provider="sglang", + mode=self.plan.mode, + ) + completed = False + try: + websocket_url = _websocket_url(self.plan.server_url, self.plan.endpoints.offer_path) + connect_started_at = time.perf_counter() + connect = self.adapter.websocket_connect_factory( + websocket_url, + open_timeout=float(self.adapter.options.connect_timeout_s), + close_timeout=float(self.adapter.options.shutdown_timeout_s), + max_size=None, + proxy=None, + ) + async with connect as websocket: + self.websocket = websocket + self.result.connected_latency_ms = (time.perf_counter() - self.started_at) * 1000.0 + self.result.offer_rtt_ms = (time.perf_counter() - connect_started_at) * 1000.0 + self.events.record("connected", websocket_url=websocket_url) + await websocket.send(msgspec.msgpack.encode(build_sglang_realtime_init(self.plan))) + self.events.record("init_sent") + await self._receive() + completed = True + except Exception as exc: # noqa: BLE001 - transport failures become results + if self._is_clean_close(exc) and self.first_frame_at is not None: + self.result.done_received = True + completed = True + self.events.record("generation_complete") + else: + self.result.error = redact_string(f"{type(exc).__name__}: {exc}") + self.events.record("session_error", error=self.result.error) + finally: + if self.control_task is not None and not self.control_task.done(): + self.control_task.cancel() + await asyncio.gather(self.control_task, return_exceptions=True) + if self.websocket is not None: + try: + await self.websocket.close() + except Exception as exc: # noqa: BLE001 - cleanup is best effort + self.events.record_error("websocket_close_failed", exc) + if completed: + self._finalize_success() + self.result.session_runtime_s = time.perf_counter() - self.started_at + event_path = await self.events.export() + self.result.artifacts_event_file = str(event_path) + return self.result + + async def _receive(self) -> None: + while True: + timeout = float(self.adapter.options.frame_timeout_s) + if self.active_started_at is not None: + remaining = self.active_started_at + float(self.plan.session_duration_s) - time.perf_counter() + if remaining <= 0: + self.events.record("session_duration_elapsed") + return + timeout = min(timeout, remaining) + try: + raw_message = await asyncio.wait_for(self.websocket.recv(), timeout=timeout) + except asyncio.TimeoutError: + if self.first_frame_at is None: + raise TimeoutError("No SGLang video frame received before the frame timeout") from None + self.events.record("session_duration_elapsed") + return + self._handle_message(raw_message) + if self.control_task is not None and self.control_task.done(): + control_error = self.control_task.exception() + if control_error is not None: + raise control_error + + def _handle_message(self, raw_message: bytes | str) -> None: + now = time.perf_counter() + if self.pending_frame_header is not None: + if not isinstance(raw_message, bytes): + raise ValueError("SGLang frame payload must be binary") + header = self.pending_frame_header + self.pending_frame_header = None + self._handle_frame_batch(header, now, payload_bytes=len(raw_message)) + return + if not isinstance(raw_message, bytes): + raise ValueError("SGLang realtime messages must be MessagePack binary frames") + message = msgspec.msgpack.decode(raw_message) + if not isinstance(message, dict): + raise ValueError("SGLang realtime message must decode to a mapping") + self._mark_metadata(now) + message_type = message.get("type") + if message_type == "frame_batch_header": + self.pending_frame_header = message + elif message_type == "frame_batch": + payload = message.get("payload") + self._handle_frame_batch( + message, + now, + payload_bytes=len(payload) if isinstance(payload, bytes) else None, + ) + elif message_type == "chunk_stats": + self._handle_chunk_stats(message, now) + elif message_type == "error": + raise RuntimeError(str(message.get("error") or message.get("message") or "SGLang realtime error")) + else: + self.events.record("sglang_message", message_type=message_type) + + def _mark_metadata(self, now: float) -> None: + self.result.metadata_messages += 1 + if self.first_metadata_at is None: + self.first_metadata_at = now + self.result.first_metadata_latency_ms = (now - self.started_at) * 1000.0 + + def _handle_frame_batch(self, message: Mapping[str, Any], now: float, *, payload_bytes: int | None) -> None: + frames = int(message.get("num_frames", 0)) + if frames <= 0: + raise ValueError("SGLang frame batch must report a positive num_frames") + request_id = message.get("request_id") + if isinstance(request_id, str) and request_id: + self.result.session_id = request_id + self.events.set_session_id(request_id) + self.result.frames_received += frames + if self.first_frame_at is None: + self.first_frame_at = now + self.first_frame_batch_size = frames + self.active_started_at = now + self.result.first_frame_latency_ms = (now - self.started_at) * 1000.0 + self.events.record("first_frame", frames=frames) + if self.plan.control_trace: + self.control_task = asyncio.create_task(self._send_control_trace()) + self.last_frame_at = now + self._mark_control_frame(message.get("event_id"), now) + self.events.record( + "frame_batch", + chunk_index=message.get("chunk_index"), + frames=frames, + payload_bytes=payload_bytes, + content_type=message.get("content_type"), + ) + + def _handle_chunk_stats(self, message: Mapping[str, Any], now: float) -> None: + self.result.status_messages += 1 + self.result.last_status_stage = "chunk_stats" + self.result.chunk_measurements.append( + StreamChunkMeasurement( + index=int(message["chunk_index"]), + frames=int(message.get("num_frames", 0)), + request_prepare_seconds=self._milliseconds(message.get("request_prepare_ms")), + compute_seconds=self._milliseconds(message.get("scheduler_forward_ms")) or 0.0, + encode_seconds=self._milliseconds(message.get("raw_payload_build_ms")), + output_pacing_seconds=self._milliseconds(message.get("pace_wait_ms")), + output_header_write_seconds=self._milliseconds(message.get("header_write_ms")), + output_payload_write_seconds=self._milliseconds(message.get("raw_write_ms")), + output_write_seconds=self._milliseconds(message.get("ws_write_ms")), + total_seconds=self._milliseconds(message.get("chunk_total_ms")), + raw_output_bytes=self._optional_int(message.get("raw_bytes")), + wire_output_bytes=self._optional_int(message.get("ws_payload_bytes")), + output_batches=self._optional_int(message.get("num_batches")), + output_content_type=str(message["content_type"]) if message.get("content_type") else None, + ) + ) + self._mark_control_ack(message.get("event_id"), now) + self.events.record( + "chunk_stats", + chunk_index=message.get("chunk_index"), + event_id=message.get("event_id"), + ) + + async def _send_control_trace(self) -> None: + active_started_at = self.active_started_at + if active_started_at is None: + return + for event_index, entry in enumerate(self.plan.control_trace): + deadline = active_started_at + float(entry["delay_s"]) + await asyncio.sleep(max(deadline - time.perf_counter(), 0.0)) + message = dict(entry["message"]) + key = str(message.get("key", "")) + action = _KEY_TO_ACTION.get(key) + if action is None: + raise ValueError(f"Unsupported SGLang control key: {key}") + operation = message.get("action") + if operation == "press": + self.held_actions.add(action) + elif operation == "release": + self.held_actions.discard(action) + else: + raise ValueError(f"Unsupported SGLang control action: {operation}") + event_id = event_index + 1 + sent_at = time.perf_counter() + self.result.control_events.append( + ControlEventResult( + index=event_index, + scheduled_delay_s=float(entry["delay_s"]), + message=message, + sent_offset_s=sent_at - active_started_at, + ) + ) + self.control_sent_at[event_index] = sent_at + self.control_by_event_id[event_id] = event_index + event = { + "type": "event", + "kind": "camera_actions", + "event_id": event_id, + "payload": { + "mode": "state", + "transitions": [ + { + "actions": sorted(self.held_actions), + "client_ts_ms": (sent_at - active_started_at) * 1000.0, + } + ], + }, + } + await self.websocket.send(msgspec.msgpack.encode(event)) + self.events.record("control_sent", event_id=event_id, actions=sorted(self.held_actions)) + + def _mark_control_ack(self, event_id: Any, now: float) -> None: + index = self.control_by_event_id.get(event_id) + if index is None: + return + control = self.result.control_events[index] + if control.ack_latency_ms is None: + control.ack_latency_ms = max((now - self.control_sent_at[index]) * 1000.0, 0.0) + + def _mark_control_frame(self, event_id: Any, now: float) -> None: + index = self.control_by_event_id.get(event_id) + if index is None: + return + control = self.result.control_events[index] + if control.next_frame_latency_ms is None: + control.next_frame_latency_ms = max((now - self.control_sent_at[index]) * 1000.0, 0.0) + + def _finalize_success(self) -> None: + self.result.success = self.first_frame_at is not None and self.result.error is None + if self.first_frame_at is None: + self.result.error = self.result.error or "No SGLang video frame received" + return + if self.last_frame_at is not None and self.last_frame_at > self.first_frame_at: + frames_after_first_batch = self.result.frames_received - self.first_frame_batch_size + self.result.stream_fps = frames_after_first_batch / (self.last_frame_at - self.first_frame_at) + + @staticmethod + def _milliseconds(value: Any) -> float | None: + return float(value) / 1000.0 if value is not None else None + + @staticmethod + def _optional_int(value: Any) -> int | None: + return int(value) if value is not None else None + + @staticmethod + def _is_clean_close(exc: Exception) -> bool: + return isinstance(exc, websockets.exceptions.ConnectionClosedOK) + + +class SGLangRealtimeAdapter(StreamTargetMetadataMixin): + """AIPerf adapter for SGLang's MessagePack realtime video endpoint.""" + + transport = "websocket" + + def __init__( + self, + *, + contract: BenchmarkContract, + config: StreamProfileConfig, + artifacts_dir: str | Path, + websocket_connect_factory: Callable[..., Any] = websockets.connect, + http_client: StreamHttpClient | None = None, + ) -> None: + self.contract = contract + self.config = config + self.options = config.transport + self.artifacts_dir = Path(artifacts_dir) + self.websocket_connect_factory = websocket_connect_factory + self.http = http_client or StreamHttpClient() + + async def check_health(self) -> None: + """Check SGLang's contract-declared health endpoint.""" + + health_path = str(self.contract.endpoint.get("health_path", "/health")) + await self.http.check_health( + f"{self.config.server_url}{health_path}", + timeout_s=float(self.options.connect_timeout_s), + ) + + async def run_session(self, plan: StreamSessionPlan) -> SessionResult: + """Execute one normalized plan through SGLang realtime video.""" + + return await _SGLangRealtimeSession(adapter=self, plan=plan).run() + + async def aclose(self) -> None: + """Close adapter-owned HTTP resources.""" + + await self.http.aclose() diff --git a/benchmarks/telefuser_aiperf/tests/test_sglang_adapter.py b/benchmarks/telefuser_aiperf/tests/test_sglang_adapter.py new file mode 100644 index 0000000..ba399f5 --- /dev/null +++ b/benchmarks/telefuser_aiperf/tests/test_sglang_adapter.py @@ -0,0 +1,264 @@ +from __future__ import annotations + +import asyncio +from pathlib import Path +from typing import Any + +import msgspec +import pytest +from aiperf.streaming.adapters import create_stream_adapter +from aiperf.streaming.config import StreamProfileConfig +from aiperf.streaming.contracts import BenchmarkContract +from aiperf.streaming.models import StreamEndpointPaths, StreamSessionPlan +from telefuser_aiperf import register_adapters +from telefuser_aiperf.payload import build_sglang_realtime_init +from telefuser_aiperf.sglang_adapter import SGLangRealtimeAdapter, _websocket_url + + +class _FakeHttpClient: + def __init__(self) -> None: + self.health_urls: list[str] = [] + + async def check_health(self, url: str, *, timeout_s: float) -> None: + self.health_urls.append(url) + + async def aclose(self) -> None: + return None + + +class _FakeWebSocket: + def __init__(self) -> None: + self.sent: list[dict[str, Any]] = [] + self.closed = False + self.messages = [ + msgspec.msgpack.encode( + { + "type": "frame_batch", + "request_id": "sglang-request", + "chunk_index": 0, + "event_id": None, + "num_frames": 13, + "content_type": "image/webp", + "payload": b"first", + } + ), + msgspec.msgpack.encode( + { + "type": "chunk_stats", + "chunk_index": 0, + "event_id": 1, + "num_frames": 13, + "request_prepare_ms": 2.0, + "scheduler_forward_ms": 500.0, + "raw_payload_build_ms": 100.0, + "pace_wait_ms": 0.0, + "header_write_ms": 1.0, + "raw_write_ms": 3.0, + "ws_write_ms": 4.0, + "chunk_total_ms": 607.0, + "raw_bytes": 1000, + "ws_payload_bytes": 600, + "num_batches": 1, + "content_type": "image/webp", + } + ), + msgspec.msgpack.encode( + { + "type": "frame_batch_header", + "request_id": "sglang-request", + "chunk_index": 1, + "event_id": 1, + "num_frames": 16, + "content_type": "image/webp", + } + ), + b"second", + ] + + async def send(self, payload: bytes) -> None: + self.sent.append(msgspec.msgpack.decode(payload)) + + async def recv(self) -> bytes: + await asyncio.sleep(0.001) + if self.messages: + return self.messages.pop(0) + await asyncio.sleep(1.0) + raise AssertionError("unreachable") + + async def close(self) -> None: + self.closed = True + + +class _FakeConnect: + def __init__(self, websocket: _FakeWebSocket) -> None: + self.websocket = websocket + + async def __aenter__(self) -> _FakeWebSocket: + return self.websocket + + async def __aexit__(self, exc_type: Any, exc: Any, traceback: Any) -> None: + return None + + +class _FakeConnectFactory: + def __init__(self) -> None: + self.websocket = _FakeWebSocket() + self.calls: list[tuple[str, dict[str, Any]]] = [] + + def __call__(self, url: str, **kwargs: Any) -> _FakeConnect: + self.calls.append((url, kwargs)) + return _FakeConnect(self.websocket) + + +def _contract() -> BenchmarkContract: + return BenchmarkContract( + contract_version="v1", + name="sglang-adapter-test", + mode="stream_world", + implementation="sglang", + model_family="lingbot_world_v2", + model="world-model", + supported_tasks=["bidirectional"], + transport="websocket", + adapter="sglang_realtime", + transport_provider="sglang", + endpoint={ + "health_path": "/health", + "offer_path": "/v1/realtime_video/generate", + }, + request_encoding={"format": "msgpack"}, + result_delivery={"media": "websocket_frame_batch"}, + workload={"size": "832x480"}, + metrics=["first_frame_latency_ms"], + artifacts={"config": "stream.json"}, + ) + + +def _config(tmp_path: Path) -> StreamProfileConfig: + return StreamProfileConfig.model_validate( + { + "contract": "contract.yaml", + "server_url": "http://127.0.0.1:30000", + "prompt": "walk forward", + "artifacts_dir": str(tmp_path), + "transport": { + "connect_timeout_s": 0.5, + "message_timeout_s": 0.5, + "frame_timeout_s": 0.5, + "shutdown_timeout_s": 0.5, + }, + } + ) + + +def _plan(image_path: Path) -> StreamSessionPlan: + return StreamSessionPlan( + logical_session_index=0, + phase="profiling", + planned_session_id="planned", + server_url="http://127.0.0.1:30000", + endpoints=StreamEndpointPaths( + health_path="/health", + offer_path="/v1/realtime_video/generate", + ), + mode="bidirectional", + task="bidirectional", + prompt="walk forward", + fps=16, + session_duration_s=0.02, + image_path=str(image_path), + request_extra={ + "size": "832x480", + "max_chunks": 60, + "num_inference_steps": 4, + }, + control_trace=[ + { + "delay_s": 0.0, + "message": { + "type": "control", + "key": "ArrowUp", + "action": "press", + }, + } + ], + ) + + +def test_registration_uses_unmodified_aiperf_registry(tmp_path: Path) -> None: + register_adapters(replace=True) + + adapter = create_stream_adapter( + contract=_contract(), + config=_config(tmp_path), + artifacts_dir=tmp_path, + ) + + assert isinstance(adapter, SGLangRealtimeAdapter) + assert adapter.transport == "websocket" + + +def test_init_payload_reads_first_frame_and_rejects_overrides(tmp_path: Path) -> None: + image_path = tmp_path / "frame.jpg" + image_path.write_bytes(b"jpeg") + plan = _plan(image_path) + + payload = build_sglang_realtime_init(plan) + + assert payload["first_frame"] == b"jpeg" + assert payload["fps"] == 16 + assert payload["max_chunks"] == 60 + with pytest.raises(ValueError, match="prompt"): + build_sglang_realtime_init(plan.model_copy(update={"request_extra": {"prompt": "override"}})) + + +def test_websocket_url_preserves_server_base_path() -> None: + assert ( + _websocket_url("https://example.test/api/", "/v1/realtime_video/generate") + == "wss://example.test/api/v1/realtime_video/generate" + ) + + +@pytest.mark.asyncio +async def test_adapter_maps_frames_controls_and_chunk_stats(tmp_path: Path) -> None: + image_path = tmp_path / "frame.jpg" + image_path.write_bytes(b"jpeg") + factory = _FakeConnectFactory() + http = _FakeHttpClient() + adapter = SGLangRealtimeAdapter( + contract=_contract(), + config=_config(tmp_path), + artifacts_dir=tmp_path, + websocket_connect_factory=factory, + http_client=http, + ) + + await adapter.check_health() + result = await adapter.run_session(_plan(image_path)) + + assert result.success is True + assert result.session_id == "sglang-request" + assert result.frames_received == 29 + assert result.stream_fps is not None + assert result.first_frame_latency_ms is not None + assert result.chunk_measurements[0].compute_seconds == 0.5 + assert result.chunk_measurements[0].encode_seconds == 0.1 + assert result.chunk_measurements[0].wire_output_bytes == 600 + assert len(result.control_events) == 1 + assert result.control_events[0].ack_latency_ms is not None + assert result.control_events[0].next_frame_latency_ms is not None + assert factory.calls[0][0] == "ws://127.0.0.1:30000/v1/realtime_video/generate" + assert factory.calls[0][1]["proxy"] is None + assert factory.websocket.sent[0]["type"] == "init" + assert factory.websocket.sent[1] == { + "type": "event", + "kind": "camera_actions", + "event_id": 1, + "payload": { + "mode": "state", + "transitions": [{"actions": ["w"], "client_ts_ms": pytest.approx(0.0, abs=20.0)}], + }, + } + assert factory.websocket.closed is True + assert http.health_urls == ["http://127.0.0.1:30000/health"] + assert Path(result.artifacts_event_file or "").is_file() diff --git a/docs/en/benchmark_aiperf.md b/docs/en/benchmark_aiperf.md index b04d320..a703129 100644 --- a/docs/en/benchmark_aiperf.md +++ b/docs/en/benchmark_aiperf.md @@ -2,7 +2,8 @@ TeleFuser exposes raw target-side facts; AIPerf owns workload execution, aggregation, resource collection, artifacts, GreptimeDB history, and visualization. The checked-in integration covers batch video generation through the -OpenAI-compatible `/v1/videos` API and LingBot streaming through LiveKit. +OpenAI-compatible `/v1/videos` API, TeleFuser LingBot streaming through LiveKit, and SGLang LingBot streaming through +its native realtime WebSocket endpoint. AIPerf's stream runner and result schema are transport-neutral. The LiveKit adapter is maintained by TeleFuser, loads from source at process startup, and produces AIPerf's standard session results. The contract records WebRTC as @@ -60,6 +61,19 @@ minutes for the command to finish. Success is `Stream profile sessions: 1/1 succ `artifacts/telefuser_aiperf/stream_lingbot_v2_1min/`. See the canonical README for model file layout, manual Python environment selection, history setup, and troubleshooting. +To profile SGLang on the same four-GPU workload instead, start its server and select the SGLang config: + +```bash +bash benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh + +bash benchmarks/telefuser_aiperf/scripts/run_stream_bench.sh \ + benchmarks/telefuser_aiperf/configs/stream_sglang_lingbot_world_v2_4gpu_1min.json +``` + +This launch explicitly uses `--flow-shift 10` for parity with the official model and TeleFuser. SGLang's source +default is `5`, so runs made with `SGLANG_FLOW_SHIFT=5` describe SGLang's default but are not model-configuration +parity comparisons. See the benchmark README for model, GPU, port, and executable overrides. + ## Ownership and metric semantics | Component | Owner | Responsibility | @@ -67,6 +81,7 @@ environment selection, history setup, and troubleshooting. | TeleFuser runtime | TeleFuser | Emit synchronized phase, chunk, runtime, cache, and environment facts | | Batch target adapter | AIPerf | Convert `/v1/videos` HTTP events into the standard request timeline | | LiveKit source adapter | TeleFuser | Convert room, track, status, metrics, and control events into session results | +| SGLang source adapter | TeleFuser | Convert MessagePack frames, chunk timings, and camera events into session results | | Aggregation and history | AIPerf | Apply warmup, percentiles, throughput, artifacts, GreptimeDB, and visualization | | Contracts and workloads | TeleFuser | Fix target capabilities, inputs, settings, and reproducible launch commands | diff --git a/scripts/setup_aiperf.sh b/scripts/setup_aiperf.sh index 2c1efce..0a8e0bc 100755 --- a/scripts/setup_aiperf.sh +++ b/scripts/setup_aiperf.sh @@ -6,12 +6,15 @@ PYTHON_BIN="${AIPERF_PYTHON_BIN:-python3}" AIPERF_ENV_DIR="${AIPERF_ENV_DIR:-${ROOT_DIR}/.venv-aiperf}" AIPERF_SPEC="${AIPERF_SPEC:-aiperf @ git+https://github.com/ActivePeter/aiperf.git@e977ffbb1648510acec431b2a3fbd1a0f7bb8a35}" LIVEKIT_SPEC="${LIVEKIT_SPEC:-livekit>=1.1.13,<2.0.0}" +MSGSPEC_SPEC="${MSGSPEC_SPEC:-msgspec>=0.18,<1.0}" +WEBSOCKETS_SPEC="${WEBSOCKETS_SPEC:-websockets>=15,<17}" ADAPTER_ROOT="${ROOT_DIR}/benchmarks/telefuser_aiperf" usage() { echo "Usage: scripts/setup_aiperf.sh" echo "" - echo "Environment overrides: AIPERF_SPEC, LIVEKIT_SPEC, AIPERF_ENV_DIR, AIPERF_PYTHON_BIN" + echo "Environment overrides: AIPERF_SPEC, LIVEKIT_SPEC, MSGSPEC_SPEC, WEBSOCKETS_SPEC," + echo " AIPERF_ENV_DIR, AIPERF_PYTHON_BIN" } if [[ $# -gt 0 ]]; then @@ -37,7 +40,11 @@ if "${AIPERF_ENV_DIR}/bin/python" -c 'import importlib.util; raise SystemExit(im "${AIPERF_ENV_DIR}/bin/python" -m pip uninstall -y aiperf fi -"${AIPERF_ENV_DIR}/bin/python" -m pip install "${AIPERF_SPEC}" "${LIVEKIT_SPEC}" +"${AIPERF_ENV_DIR}/bin/python" -m pip install \ + "${AIPERF_SPEC}" \ + "${LIVEKIT_SPEC}" \ + "${MSGSPEC_SPEC}" \ + "${WEBSOCKETS_SPEC}" mkdir -p "${ROOT_DIR}/artifacts" From 9dc620daaf3c7ff98130c0410e304e2a004ddc03 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Sun, 2 Aug 2026 06:26:24 +0000 Subject: [PATCH 02/11] perf(lingbot): optimize streaming DiT execution Fuse normalization and modulation kernels, reuse LingBot model inputs and projected session data, and reduce redundant RoPE and KV bookkeeping across denoising steps. Keep FA4 return-LSE dispatch compatible and prefer tf-kernel RMSNorm when available. Add focused regression coverage for the optimized ops, cache invalidation, and denoising lifecycle. --- telefuser/kernel/triton/__init__.py | 4 + telefuser/kernel/triton/scale_shift.py | 148 ++++++++++ telefuser/models/lingbot_world_fast_dit.py | 245 ++++++++++++---- telefuser/ops/__init__.py | 8 +- telefuser/ops/attention/attention_impl.py | 8 +- telefuser/ops/normalization.py | 70 ++++- .../pipelines/lingbot_world_fast/denoising.py | 273 ++++++++++++++++-- .../models/test_lingbot_world_fast_dit.py | 243 ++++++++++++++++ tests/unit/ops/test_attention_backends.py | 60 +++- tests/unit/ops/test_normalization.py | 33 ++- .../lingbot_world_fast/test_parallelism.py | 1 + .../test_runtime_baseline.py | 19 ++ .../lingbot_world_fast/test_session_cache.py | 76 ++++- 13 files changed, 1097 insertions(+), 91 deletions(-) diff --git a/telefuser/kernel/triton/__init__.py b/telefuser/kernel/triton/__init__.py index 0a0bd85..80d7e60 100644 --- a/telefuser/kernel/triton/__init__.py +++ b/telefuser/kernel/triton/__init__.py @@ -20,6 +20,8 @@ from .quant import per_token_dequant_fp8, per_token_quant_fp8 from .rotary import apply_rotary_embedding from .scale_shift import ( + fused_add_layernorm_scale_shift, + fused_layernorm_scale_shift, fused_layernorm_scale_shift_gate_select01, fused_residual_layernorm_scale_shift_gate_select01, fused_scale_shift, @@ -34,6 +36,8 @@ "apply_rotary_embedding", "apply_rotary_embedding_inplace", "fused_scale_shift", + "fused_layernorm_scale_shift", + "fused_add_layernorm_scale_shift", "fused_scale_shift_gate_select", "fused_layernorm_scale_shift_gate_select01", "fused_residual_layernorm_scale_shift_gate_select01", diff --git a/telefuser/kernel/triton/scale_shift.py b/telefuser/kernel/triton/scale_shift.py index 18c1115..dd50b48 100644 --- a/telefuser/kernel/triton/scale_shift.py +++ b/telefuser/kernel/triton/scale_shift.py @@ -1062,3 +1062,151 @@ def fused_scale_shift_gate_select( num_stages=2, ) return output, gate_out + + +@triton.jit +def _fused_layernorm_scale_shift_kernel( + output_ptr, + residual_out_ptr, + x_ptr, + residual_ptr, + weight_ptr, + bias_ptr, + scale_ptr, + shift_ptr, + inner_dim, + seq_len, + scale_seq_len, + shift_seq_len, + eps, + HAS_RESIDUAL: tl.constexpr, + STORE_RESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + BLOCK_N: tl.constexpr, +): + row = tl.program_id(0) + cols = tl.arange(0, BLOCK_N) + mask = cols < inner_dim + batch_idx = row // seq_len + seq_idx = row % seq_len + offsets = row * inner_dim + cols + + value = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + if HAS_RESIDUAL: + value += tl.load(residual_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + if STORE_RESIDUAL: + tl.store(residual_out_ptr + offsets, value, mask=mask) + + normalized = _compute_layernorm( + value, + cols, + mask, + inner_dim, + eps, + weight_ptr, + bias_ptr, + 1, + 1, + HAS_WEIGHT, + HAS_BIAS, + ) + scale_row = batch_idx * scale_seq_len + tl.minimum(seq_idx, scale_seq_len - 1) + shift_row = batch_idx * shift_seq_len + tl.minimum(seq_idx, shift_seq_len - 1) + scale = tl.load(scale_ptr + scale_row * inner_dim + cols, mask=mask, other=0.0).to(tl.float32) + shift = tl.load(shift_ptr + shift_row * inner_dim + cols, mask=mask, other=0.0).to(tl.float32) + tl.store(output_ptr + offsets, normalized * (1.0 + scale) + shift, mask=mask) + + +def _normalize_modulation(tensor: torch.Tensor, batch_size: int, seq_len: int, hidden_size: int) -> torch.Tensor: + if tensor.dim() == 2: + tensor = tensor.unsqueeze(1) + if tensor.dim() != 3 or tensor.shape[0] not in (1, batch_size) or tensor.shape[1] not in (1, seq_len): + raise ValueError("scale and shift must broadcast to [B, L, C]") + if tensor.shape[2] != hidden_size: + raise ValueError(f"scale and shift hidden size must be {hidden_size}, got {tensor.shape[2]}") + return tensor.expand(batch_size, tensor.shape[1], hidden_size).contiguous() + + +def fused_layernorm_scale_shift( + x: torch.Tensor, + weight: torch.Tensor | None, + bias: torch.Tensor | None, + scale: torch.Tensor, + shift: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Fuse LayerNorm and adaptive scale/shift for contiguous BLC tensors.""" + assert x.is_cuda and x.is_contiguous() + batch_size, seq_len, hidden_size = x.shape + scale = _normalize_modulation(scale, batch_size, seq_len, hidden_size) + shift = _normalize_modulation(shift, batch_size, seq_len, hidden_size) + output = torch.empty_like(x) + placeholder = x + block_n = triton.next_power_of_2(hidden_size) + _fused_layernorm_scale_shift_kernel[(batch_size * seq_len,)]( + output, + placeholder, + x, + placeholder, + placeholder if weight is None else weight.contiguous(), + placeholder if bias is None else bias.contiguous(), + scale, + shift, + hidden_size, + seq_len, + scale.shape[1], + shift.shape[1], + eps, + HAS_RESIDUAL=False, + STORE_RESIDUAL=False, + HAS_WEIGHT=weight is not None, + HAS_BIAS=bias is not None, + BLOCK_N=block_n, + num_warps=8, + ) + return output + + +def fused_add_layernorm_scale_shift( + residual: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor | None, + bias: torch.Tensor | None, + scale: torch.Tensor, + shift: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fuse residual add, LayerNorm, and adaptive scale/shift for BLC tensors.""" + assert residual.is_cuda and residual.is_contiguous() and x.is_contiguous() + if residual.shape != x.shape: + raise ValueError(f"residual and x shapes must match, got {residual.shape} and {x.shape}") + batch_size, seq_len, hidden_size = x.shape + scale = _normalize_modulation(scale, batch_size, seq_len, hidden_size) + shift = _normalize_modulation(shift, batch_size, seq_len, hidden_size) + output = torch.empty_like(x) + residual_out = torch.empty_like(residual) + placeholder = x + block_n = triton.next_power_of_2(hidden_size) + _fused_layernorm_scale_shift_kernel[(batch_size * seq_len,)]( + output, + residual_out, + x, + residual, + placeholder if weight is None else weight.contiguous(), + placeholder if bias is None else bias.contiguous(), + scale, + shift, + hidden_size, + seq_len, + scale.shape[1], + shift.shape[1], + eps, + HAS_RESIDUAL=True, + STORE_RESIDUAL=True, + HAS_WEIGHT=weight is not None, + HAS_BIAS=bias is not None, + BLOCK_N=block_n, + num_warps=8, + ) + return output, residual_out diff --git a/telefuser/models/lingbot_world_fast_dit.py b/telefuser/models/lingbot_world_fast_dit.py index 37a8bee..a6360bb 100644 --- a/telefuser/models/lingbot_world_fast_dit.py +++ b/telefuser/models/lingbot_world_fast_dit.py @@ -16,21 +16,40 @@ from telefuser.distributed.parallel_shard import sequence_parallel_shard, sequence_parallel_unshard from telefuser.distributed.ulysses_comm import ulysses_gather_heads, ulysses_scatter_heads from telefuser.ops.attention import attention as attn_func -from telefuser.ops.normalization import LayerNorm, RMSNorm +from telefuser.ops.normalization import ( + LayerNorm, + RMSNorm, + _fused_add_layer_norm_scale_shift, + _fused_layer_norm_scale_shift, + fused_scale_shift, +) from telefuser.utils.logging import logger from telefuser.utils.model_weight import init_weights_on_device, load_state_dict from .wan_video_dit import apply_rotary_emb, precompute_freqs_cis_3d, sinusoidal_embedding_1d +_PreparedControl = tuple[torch.Tensor, tuple[tuple[torch.Tensor, torch.Tensor], ...]] -def _cache_index_to_int(value: int | torch.Tensor) -> int: + +def _cache_index_to_int(cache: dict[str, object], name: str) -> int: + host_indices = cache.get("host_indices") + if isinstance(host_indices, dict): + value = host_indices.get(name) + if isinstance(value, int): + return value + value = cache[name] if isinstance(value, int): return value - return int(value.item()) + if not isinstance(value, torch.Tensor): + raise TypeError(f"Unsupported LingBot cache index type: {type(value).__name__}") + resolved = int(value.item()) + cache.setdefault("host_indices", {})[name] = resolved + return resolved -def _set_cache_index(cache: dict[str, torch.Tensor | int], name: str, value: int) -> None: +def _set_cache_index(cache: dict[str, object], name: str, value: int) -> None: """Update an eager cursor in place when its storage is a scalar tensor.""" + cache.setdefault("host_indices", {})[name] = value cursor = cache[name] if isinstance(cursor, torch.Tensor): cursor.fill_(value) @@ -68,14 +87,13 @@ def __init__( self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() - def _apply_causal_rope( + def _prepare_causal_rope( self, - x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor, grid_size: tuple[int, int, int], start_frame: int, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor, torch.Tensor]: f, h, w = grid_size head_dim = self.head_dim @@ -104,10 +122,25 @@ def _apply_causal_rope( dim=-1, ).reshape(seq_len, 1, -1) - roped = apply_rotary_emb(x[:, :seq_len], (cos, sin)) - if x.shape[1] == seq_len: + return cos, sin + + def _apply_causal_rope( + self, + x: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + grid_size: tuple[int, int, int], + start_frame: int, + causal_rope: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + if causal_rope is None: + causal_rope = self._prepare_causal_rope(freqs_cos, freqs_sin, grid_size, start_frame) + + sequence_length = math.prod(grid_size) + roped = apply_rotary_emb(x[:, :sequence_length], causal_rope) + if x.shape[1] == sequence_length: return roped - return torch.cat([roped, x[:, seq_len:]], dim=1) + return torch.cat([roped, x[:, sequence_length:]], dim=1) def forward( self, @@ -119,29 +152,27 @@ def forward( current_start: int, max_attention_size: int, device_mesh: DeviceMesh | None = None, - ) -> torch.Tensor: + causal_rope: tuple[torch.Tensor, torch.Tensor] | None = None, + update_cache_only: bool = False, + ) -> torch.Tensor | None: group = get_ulysses_group(device_mesh) ulysses_enabled = group is not None and get_ulysses_world_size(device_mesh) > 1 q = rearrange(self.norm_q(self.q(x)), "b s (n d) -> b s n d", n=self.num_heads) - q_wait = ulysses_scatter_heads(q, group) if ulysses_enabled else None k = rearrange(self.norm_k(self.k(x)), "b s (n d) -> b s n d", n=self.num_heads) - k_wait = ulysses_scatter_heads(k, group) if ulysses_enabled else None v = rearrange(self.v(x), "b s (n d) -> b s n d", n=self.num_heads) - v_wait = ulysses_scatter_heads(v, group) if ulysses_enabled else None + qkv_wait = ulysses_scatter_heads(torch.cat((q, k, v), dim=-1), group) if ulysses_enabled else None frame_tokens = grid_size[1] * grid_size[2] start_frame = current_start // frame_tokens valid_seq_len = math.prod(grid_size) if ulysses_enabled: - q = q_wait() - k = k_wait() - v = v_wait() + q, k, v = qkv_wait().chunk(3, dim=-1) padded_seq_len = q.shape[1] - q = self._apply_causal_rope(q, freqs_cos, freqs_sin, grid_size, start_frame)[:, :valid_seq_len] - k = self._apply_causal_rope(k, freqs_cos, freqs_sin, grid_size, start_frame)[:, :valid_seq_len] + q = self._apply_causal_rope(q, freqs_cos, freqs_sin, grid_size, start_frame, causal_rope)[:, :valid_seq_len] + k = self._apply_causal_rope(k, freqs_cos, freqs_sin, grid_size, start_frame, causal_rope)[:, :valid_seq_len] v = v[:, :valid_seq_len] else: - q = self._apply_causal_rope(q, freqs_cos, freqs_sin, grid_size, start_frame) - k = self._apply_causal_rope(k, freqs_cos, freqs_sin, grid_size, start_frame) + q = self._apply_causal_rope(q, freqs_cos, freqs_sin, grid_size, start_frame, causal_rope) + k = self._apply_causal_rope(k, freqs_cos, freqs_sin, grid_size, start_frame, causal_rope) num_new_tokens = q.shape[1] current_end = current_start + num_new_tokens @@ -150,8 +181,8 @@ def forward( cache_k = kv_cache["k"] cache_v = kv_cache["v"] kv_cache_size = cache_k.shape[1] - global_end = _cache_index_to_int(kv_cache["global_end_index"]) - local_end = _cache_index_to_int(kv_cache["local_end_index"]) + global_end = _cache_index_to_int(kv_cache, "global_end_index") + local_end = _cache_index_to_int(kv_cache, "local_end_index") if self.local_attn_size != -1 and current_end > global_end and num_new_tokens + local_end > kv_cache_size: evicted = num_new_tokens + local_end - kv_cache_size @@ -170,6 +201,11 @@ def forward( cache_k[:, local_start:local_end] = k cache_v[:, local_start:local_end] = v + _set_cache_index(kv_cache, "global_end_index", current_end) + _set_cache_index(kv_cache, "local_end_index", local_end) + if update_cache_only: + return None + attn_start = max(0, local_end - max_attention_size) k_cache = cache_k[:, attn_start:local_end] v_cache = cache_v[:, attn_start:local_end] @@ -183,9 +219,6 @@ def forward( output_layout="BSND", ) - _set_cache_index(kv_cache, "global_end_index", current_end) - _set_cache_index(kv_cache, "local_end_index", local_end) - if ulysses_enabled: pad_len = padded_seq_len - num_new_tokens if pad_len > 0: @@ -263,7 +296,7 @@ def forward( class Gate(nn.Module): def forward(self, x: torch.Tensor, gate: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: - return x + gate * residual + return fused_scale_shift(residual, gate, x, scale_constant=0.0).to(dtype=x.dtype) class LingBotWorldFastBlock(nn.Module): @@ -319,21 +352,21 @@ def forward( current_start: int, max_attention_size: int, control_tokens: torch.Tensor | None = None, + camera_modulation: tuple[torch.Tensor, torch.Tensor] | None = None, device_mesh: DeviceMesh | None = None, + causal_rope: tuple[torch.Tensor, torch.Tensor] | None = None, + update_cache_only: bool = False, ) -> torch.Tensor: modulation = self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (modulation.unsqueeze(0) + t_mod).chunk( - 6, dim=2 - ) - - shift_msa = shift_msa.squeeze(2) - scale_msa = scale_msa.squeeze(2) - gate_msa = gate_msa.squeeze(2) - shift_mlp = shift_mlp.squeeze(2) - scale_mlp = scale_mlp.squeeze(2) - gate_mlp = gate_mlp.squeeze(2) + if t_mod.dim() == 4: + modulation_chunks = (modulation.unsqueeze(0) + t_mod).chunk(6, dim=2) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + chunk.squeeze(2) for chunk in modulation_chunks + ) + else: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (modulation + t_mod).chunk(6, dim=1) - attn_in = self.norm1(x) * (1 + scale_msa) + shift_msa + attn_in = _fused_layer_norm_scale_shift(x, scale_msa, shift_msa, eps=self.norm1.eps) attn_out = self.self_attn( attn_in, freqs_cos=freqs_cos, @@ -343,16 +376,31 @@ def forward( current_start=current_start, max_attention_size=max_attention_size, device_mesh=device_mesh, + causal_rope=causal_rope, + update_cache_only=update_cache_only, ) + if update_cache_only: + return x + if attn_out is None: + raise RuntimeError("LingBot self-attention returned no output during denoising") x = self.gate(x, gate_msa, attn_out) - if control_tokens is not None: + if camera_modulation is not None: + scale, shift = camera_modulation + x = fused_scale_shift(x, scale, shift).to(dtype=x.dtype) + elif control_tokens is not None: hidden = self.cam_injector_layer2(F.silu(self.cam_injector_layer1(control_tokens))) hidden = hidden + control_tokens x = (1.0 + self.cam_scale_layer(hidden)) * x + self.cam_shift_layer(hidden) - x = x + self.cross_attn(self.norm3(x), context, crossattn_cache) - mlp_in = self.norm2(x) * (1 + scale_mlp) + shift_mlp + cross_attn_out = self.cross_attn(self.norm3(x), context, crossattn_cache) + mlp_in, x = _fused_add_layer_norm_scale_shift( + x, + cross_attn_out, + scale_mlp, + shift_mlp, + eps=self.norm2.eps, + ) x = self.gate(x, gate_mlp, self.ffn(mlp_in)) return x @@ -372,7 +420,8 @@ def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: shift, scale = (modulation.unsqueeze(0) + t.unsqueeze(2)).chunk(2, dim=2) shift = shift.squeeze(2) scale = scale.squeeze(2) - return self.head(self.norm(x) * (1 + scale) + shift) + modulated = _fused_layer_norm_scale_shift(x, scale, shift, eps=self.norm.eps) + return self.head(modulated) class LingBotWorldFastDiT(BaseModel): @@ -480,7 +529,7 @@ def _build_timestep_embeddings(self, t: torch.Tensor, seq_len: int) -> tuple[tor if t.dim() == 1: emb_input = sinusoidal_embedding_1d(self.freq_dim, t).float() emb = self.time_embedding(emb_input) - t_mod = self.time_projection(emb).unflatten(1, (6, self.dim)).unsqueeze(1).expand(-1, seq_len, -1, -1) + t_mod = self.time_projection(emb).unflatten(1, (6, self.dim)) return emb.unsqueeze(1).float(), t_mod.float() flat = t.flatten() @@ -490,7 +539,13 @@ def _build_timestep_embeddings(self, t: torch.Tensor, seq_len: int) -> tuple[tor t_mod = self.time_projection(emb).unflatten(2, (6, self.dim)) return emb.float(), t_mod.float() - def _prepare_control_tokens(self, control_tensor: torch.Tensor | None) -> torch.Tensor | None: + def _prepare_control( + self, + control_tensor: torch.Tensor | None, + *, + shard_for_usp: bool = False, + ) -> _PreparedControl | None: + """Prepare the chunk-invariant camera tensors consumed by every block.""" if control_tensor is None: return None @@ -507,7 +562,46 @@ def _prepare_control_tokens(self, control_tensor: torch.Tensor | None) -> torch. ) control_tokens = self.patch_embedding_wancamctrl(control_tokens) hidden = self.c2ws_hidden_states_layer2(F.silu(self.c2ws_hidden_states_layer1(control_tokens))) - return control_tokens + hidden + control_tokens = control_tokens + hidden + if shard_for_usp: + sequence_parallel_shard(self.device_mesh, [control_tokens], [1]) + camera_modulations: list[tuple[torch.Tensor, torch.Tensor]] = [] + for block in self.blocks: + block_hidden = block.cam_injector_layer2(F.silu(block.cam_injector_layer1(control_tokens))) + block_hidden = block_hidden + control_tokens + camera_modulations.append((block.cam_scale_layer(block_hidden), block.cam_shift_layer(block_hidden))) + return control_tokens, tuple(camera_modulations) + + def _project_text_context(self, context: torch.Tensor) -> torch.Tensor: + """Project static prompt embeddings once before repeated chunk forwards.""" + return self.text_embedding(context) + + def _prepare_session_causal_rope( + self, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + grid_size: tuple[int, int, int], + current_start: int, + session_input_cache: dict[str, object] | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Reuse one chunk's 3D RoPE across every block and DiT forward.""" + key = (grid_size, current_start, freqs_cos.device, freqs_cos.dtype) + cached = None if session_input_cache is None else session_input_cache.get("causal_rope") + if isinstance(cached, tuple) and len(cached) == 2 and cached[0] == key: + value = cached[1] + if isinstance(value, tuple) and len(value) == 2: + return value + + frame_tokens = grid_size[1] * grid_size[2] + rope = self.blocks[0].self_attn._prepare_causal_rope( + freqs_cos, + freqs_sin, + grid_size, + current_start // frame_tokens, + ) + if session_input_cache is not None: + session_input_cache["causal_rope"] = (key, rope) + return rope def patchify(self, x: torch.Tensor) -> tuple[torch.Tensor, tuple[int, int, int]]: x = x.contiguous(memory_format=torch.channels_last_3d) @@ -538,7 +632,13 @@ def forward( crossattn_cache: list[dict[str, torch.Tensor | bool | int]] | None = None, current_start: int = 0, max_attention_size: int = 1_000_000, - ) -> torch.Tensor: + *, + _session_input_cache: dict[str, object] | None = None, + _prepared_control: _PreparedControl | None = None, + _prepared_control_is_sharded: bool = False, + _projected_context: torch.Tensor | None = None, + update_cache_only: bool = False, + ) -> torch.Tensor | None: if y is not None: y = y.to(device=x.device, dtype=x.dtype) x = torch.cat([x, y], dim=1) @@ -546,16 +646,56 @@ def forward( x, grid_size = self.patchify(x) seq_len = x.shape[1] t_head, t_mod = self._build_timestep_embeddings(timestep, seq_len) - context = self.text_embedding(context) - control_tokens = self._prepare_control_tokens(control_tensor) + if _projected_context is None: + cached_context = None if _session_input_cache is None else _session_input_cache.get("projected_context") + if isinstance(cached_context, torch.Tensor): + context = cached_context + else: + context = self._project_text_context(context) + if _session_input_cache is not None: + _session_input_cache["projected_context"] = context + else: + context = _projected_context + if _prepared_control is None: + cached_control = None if _session_input_cache is None else _session_input_cache.get("prepared_control") + if cached_control is None: + prepared_control = self._prepare_control(control_tensor, shard_for_usp=self.usp_flag) + if _session_input_cache is not None: + _session_input_cache["prepared_control"] = prepared_control + prepared_control_is_sharded = self.usp_flag and prepared_control is not None + else: + prepared_control = cached_control + prepared_control_is_sharded = _prepared_control_is_sharded + else: + prepared_control = _prepared_control + prepared_control_is_sharded = _prepared_control_is_sharded + control_tokens = prepared_control[0] if prepared_control is not None else None + camera_modulations = prepared_control[1] if prepared_control is not None else None freqs_cos = self.freqs_cos.to(device=x.device) freqs_sin = self.freqs_sin.to(device=x.device) + causal_rope = self._prepare_session_causal_rope( + freqs_cos, + freqs_sin, + grid_size, + current_start, + _session_input_cache, + ) full_seq_len = seq_len if self.usp_flag: - shard_tensors = [x, t_mod, control_tokens] - shard_dims = [1, 1, 1] + shard_tensors = [x] + shard_dims = [1] + if t_mod.dim() == 4 and t_mod.shape[1] == seq_len: + shard_tensors.append(t_mod) + shard_dims.append(1) + if control_tokens is not None and not prepared_control_is_sharded: + shard_tensors.append(control_tokens) + shard_dims.append(1) + if camera_modulations is not None and not prepared_control_is_sharded: + for scale, shift in camera_modulations: + shard_tensors.extend((scale, shift)) + shard_dims.extend((1, 1)) if t_head.shape[1] == seq_len: shard_tensors.append(t_head) shard_dims.append(1) @@ -577,9 +717,14 @@ def forward( current_start=current_start, max_attention_size=max_attention_size, control_tokens=control_tokens, + camera_modulation=None if camera_modulations is None else camera_modulations[idx], device_mesh=self.device_mesh, + causal_rope=causal_rope, + update_cache_only=update_cache_only and idx == len(self.blocks) - 1, ) + if update_cache_only: + return None x = self.head(x, t_head) if self.usp_flag: (x,) = sequence_parallel_unshard(self.device_mesh, [x], [1], [full_seq_len]) diff --git a/telefuser/ops/__init__.py b/telefuser/ops/__init__.py index 2b76b31..8212542 100644 --- a/telefuser/ops/__init__.py +++ b/telefuser/ops/__init__.py @@ -13,7 +13,13 @@ from .base import CustomOp, CustomOpFunction from .custom_op import TritonKernelWrapper, register_custom_op from .moe import grouped_expert_forward, route_topk -from .normalization import AdaLayerNormContinuous, LayerNorm, RMSNorm, fused_scale_shift, modulate +from .normalization import ( + AdaLayerNormContinuous, + LayerNorm, + RMSNorm, + fused_scale_shift, + modulate, +) from .rotary import apply_rotary_emb __all__ = [ diff --git a/telefuser/ops/attention/attention_impl.py b/telefuser/ops/attention/attention_impl.py index 5093d71..f05b495 100755 --- a/telefuser/ops/attention/attention_impl.py +++ b/telefuser/ops/attention/attention_impl.py @@ -212,19 +212,19 @@ def attention( # Flash Attention implementations if attn_impl == AttnImplType.FLASH_ATTN_4 and FLASH_ATTN_4_AVAILABLE and flash_attn4 is not None: - result = flash_attn4(q, k, v, softmax_scale=scale, return_softmax_lse=return_lse, **kwargs) - if return_lse: + result = flash_attn4(q, k, v, softmax_scale=scale, causal=is_causal, return_lse=return_lse, **kwargs) + if isinstance(result, tuple): output, lse = result else: output = result elif attn_impl == AttnImplType.FLASH_ATTN_3 and FLASH_ATTN_3_AVAILABLE and flash_attn3 is not None: - result = flash_attn3(q, k, v, softmax_scale=scale, return_softmax_lse=return_lse, **kwargs) + result = flash_attn3(q, k, v, softmax_scale=scale, causal=is_causal, return_softmax_lse=return_lse, **kwargs) if return_lse: output, lse = result else: output = result elif attn_impl == AttnImplType.FLASH_ATTN_2 and FLASH_ATTN_2_AVAILABLE and flash_attn2 is not None: - result = flash_attn2(q, k, v, softmax_scale=scale, return_attn_probs=return_lse, **kwargs) + result = flash_attn2(q, k, v, softmax_scale=scale, causal=is_causal, return_attn_probs=return_lse, **kwargs) if return_lse: output, lse, _ = result else: diff --git a/telefuser/ops/normalization.py b/telefuser/ops/normalization.py index 38c0ef0..ae93e87 100644 --- a/telefuser/ops/normalization.py +++ b/telefuser/ops/normalization.py @@ -15,7 +15,19 @@ from .base import CustomOp -KernelName = Literal["norm_infer", "layer_norm_fn", "fused_scale_shift"] +_compiled_rmsnorm: Callable | None = None +try: + from tf_kernel import rmsnorm as _compiled_rmsnorm +except ImportError: + pass + +KernelName = Literal[ + "norm_infer", + "layer_norm_fn", + "fused_scale_shift", + "fused_layernorm_scale_shift", + "fused_add_layernorm_scale_shift", +] @functools.lru_cache(maxsize=None) @@ -33,6 +45,14 @@ def _get_triton_kernel(name: KernelName) -> Callable: from telefuser.kernel.triton import fused_scale_shift return fused_scale_shift + elif name == "fused_layernorm_scale_shift": + from telefuser.kernel.triton import fused_layernorm_scale_shift + + return fused_layernorm_scale_shift + elif name == "fused_add_layernorm_scale_shift": + from telefuser.kernel.triton import fused_add_layernorm_scale_shift + + return fused_add_layernorm_scale_shift raise ValueError(f"Unknown kernel: {name}") @@ -88,6 +108,18 @@ def forward_cuda(self, hidden_states: torch.Tensor) -> torch.Tensor: assert self.weight is not None # Ensure input is contiguous for Triton kernel hidden_states = hidden_states.contiguous() + if ( + _compiled_rmsnorm is not None + and hidden_states.dtype in (torch.float16, torch.bfloat16) + and self.weight.dtype == hidden_states.dtype + ): + shape = hidden_states.shape + normalized = _compiled_rmsnorm( + hidden_states.view(-1, shape[-1]), + self.weight, + self.eps, + ) + return normalized.view(shape) norm_infer = _get_triton_kernel("norm_infer") # weight is nn.Parameter - always contiguous by PyTorch convention return norm_infer(hidden_states, self.weight, None, self.eps, is_rms_norm=True) @@ -265,6 +297,42 @@ def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch return fused_scale_shift(x, scale, shift, scale_constant=1.0) +def _fused_layer_norm_scale_shift( + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + *, + weight: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + eps: float = 1e-6, +) -> torch.Tensor: + """Fuse LayerNorm and adaptive modulation while preserving the input dtype.""" + if torch.compiler.is_compiling() or x.device.type != "cuda": + normalized = F.layer_norm(x, (x.shape[-1],), weight, bias, eps) + return (normalized * (1 + scale) + shift).to(dtype=x.dtype) + kernel = _get_triton_kernel("fused_layernorm_scale_shift") + return kernel(x.contiguous(), weight, bias, scale, shift, eps) + + +def _fused_add_layer_norm_scale_shift( + residual: torch.Tensor, + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + *, + weight: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + eps: float = 1e-6, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fuse residual add, LayerNorm, and adaptive modulation.""" + if torch.compiler.is_compiling() or x.device.type != "cuda": + residual_out = (residual + x).to(dtype=residual.dtype) + normalized = F.layer_norm(residual_out, (residual_out.shape[-1],), weight, bias, eps) + return (normalized * (1 + scale) + shift).to(dtype=residual.dtype), residual_out + kernel = _get_triton_kernel("fused_add_layernorm_scale_shift") + return kernel(residual.contiguous(), x.contiguous(), weight, bias, scale, shift, eps) + + __all__ = [ "RMSNorm", "LayerNorm", diff --git a/telefuser/pipelines/lingbot_world_fast/denoising.py b/telefuser/pipelines/lingbot_world_fast/denoising.py index 58f33ea..de3900d 100644 --- a/telefuser/pipelines/lingbot_world_fast/denoising.py +++ b/telefuser/pipelines/lingbot_world_fast/denoising.py @@ -1,6 +1,8 @@ from __future__ import annotations -from dataclasses import dataclass +from collections import deque +from dataclasses import dataclass, field +from typing import Callable import torch @@ -14,6 +16,8 @@ from telefuser.schedulers.unipc import FlowUniPCMultistepScheduler from telefuser.utils.logging import logger +from .vae_stage import LingBotWorldFastVAEDecodeStage + def _select_timesteps( scheduler: FlowUniPCMultistepScheduler, @@ -45,7 +49,12 @@ class _DenoisingCacheState: generator: torch.Generator noise_generator: torch.Generator noise_shape: tuple[int, int, int, int, int] + prompt_emb: torch.Tensor | None = None pool_slot: int | None = None + projected_context_key: tuple[object, ...] | None = None + prepared_control_key: tuple[object, ...] | None = None + prepared_control_is_sharded: bool = False + session_input_cache: dict[str, object] = field(default_factory=dict) @dataclass(frozen=True) @@ -129,6 +138,7 @@ def try_acquire( "v": self.self_v[slot, layer], "global_end_index": self.cursors[slot, layer, 0], "local_end_index": self.cursors[slot, layer, 1], + "host_indices": {"global_end_index": 0, "local_end_index": 0}, } for layer in range(self.num_layers) ] @@ -162,8 +172,14 @@ def __init__( raise ValueError("LingBot denoising stage requires a loaded lingbot_world_fast_dit module") self.dit.set_attention_config(model_runtime_config.attention_config) self.model_names = ["dit"] + # This worker alternates persistent DiT and VAE decode calls. Keeping the + # allocator cache avoids driver allocations between chunks; cached blocks + # remain reclaimable and do not change model or session-cache capacity. + self.empty_cache_after_call = False self._cache_registry: dict[int, _DenoisingCacheState] = {} self._cache_pool: _DenoisingCachePool | None = None + self._vae_decode_stage: LingBotWorldFastVAEDecodeStage | None = None + self._pending_vae_decode_latents: dict[int, deque[torch.Tensor]] = {} if model_runtime_config.parallel_config.world_size == 1 and model_runtime_config.compile_config.enabled: logger.info(f"Enabling torch.compile for {self.name}") self.dit = torch.compile(self.dit, **model_runtime_config.compile_config.get_compile_kwargs()) @@ -189,6 +205,66 @@ def parallel_models(self) -> None: if self.model_runtime_config.compile_config.enabled: logger.info(f"Enabling torch.compile for {self.name}") self.dit = torch.compile(self.dit, **self.model_runtime_config.compile_config.get_compile_kwargs()) + if self._vae_decode_stage is not None: + self._vae_decode_stage.device = self.device + self._vae_decode_stage.model_runtime_config.device_id = self.device.index or 0 + self._vae_decode_stage.parallel_models() + + def attach_vae_decode_stage(self, stage: LingBotWorldFastVAEDecodeStage) -> None: + """Attach a VAE decoder that shares this worker's process group and CUDA context.""" + if self._vae_decode_stage is not None: + raise RuntimeError("a VAE decode stage is already attached") + self._vae_decode_stage = stage + + def _require_vae_decode_stage(self) -> LingBotWorldFastVAEDecodeStage: + if self._vae_decode_stage is None: + raise RuntimeError("no VAE decode stage is attached") + return self._vae_decode_stage + + def reset_vae_decode_device_memory_peak(self) -> bool: + return self._require_vae_decode_stage().reset_device_memory_peak() + + def vae_decode_device_memory_snapshots(self) -> list[dict[str, int | str]]: + return self._require_vae_decode_stage().device_memory_snapshots() + + def estimate_vae_decode_session_cache_bytes(self) -> int: + return self._require_vae_decode_stage().estimate_session_cache_bytes() + + def observed_vae_decode_session_cache_bytes(self) -> int: + return self._require_vae_decode_stage().observed_session_cache_bytes() + + def configure_vae_decode_cache_pool(self, capacity: int): + return self._require_vae_decode_stage().configure_cache_pool(capacity) + + def initialize_vae_decode_cache(self, cache_handle: int) -> bool: + return self._require_vae_decode_stage().initialize_cache(cache_handle) + + def decode_chunk( + self, + cache_handle: int, + latents: torch.Tensor | None, + is_first_clip: bool, + is_last_clip: bool, + _benchmark_profile: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, dict[str, float]]: + pending_latents = self._pending_vae_decode_latents.get(cache_handle) + local_latents = pending_latents.popleft() if pending_latents else None + decode_latents = local_latents if local_latents is not None else latents + if decode_latents is None: + raise RuntimeError(f"No worker-local VAE latent is available for cache handle {cache_handle}") + if pending_latents is not None and not pending_latents: + self._pending_vae_decode_latents.pop(cache_handle) + profile_kwargs = {"_benchmark_profile": True} if _benchmark_profile else {} + return self._require_vae_decode_stage().decode_chunk( + cache_handle, + decode_latents, + is_first_clip, + is_last_clip, + **profile_kwargs, + ) + + def release_vae_decode_cache(self, cache_handle: int) -> bool: + return self._require_vae_decode_stage().release_cache(cache_handle) def _init_self_kv_cache( self, @@ -217,6 +293,7 @@ def _init_self_kv_cache( # the updates made by the attention block. "global_end_index": torch.zeros((), dtype=torch.int64, device=self.device), "local_end_index": torch.zeros((), dtype=torch.int64, device=self.device), + "host_indices": {"global_end_index": 0, "local_end_index": 0}, } for _ in range(self.dit.num_layers) ] @@ -319,6 +396,7 @@ def initialize_cache( generator_state: list[int], noise_generator_state: list[int], noise_shape: tuple[int, int, int, int, int], + prompt_emb: torch.Tensor | None = None, timestep_indices: tuple[int, ...] = (0, 179, 358, 679), ) -> bool: """Atomically register session-scoped KV, scheduler, and RNG state.""" @@ -360,6 +438,7 @@ def initialize_cache( generator=generator, noise_generator=noise_generator, noise_shape=noise_shape, + prompt_emb=prompt_emb, pool_slot=pool_slot, ) except Exception: @@ -388,6 +467,28 @@ def _convert_flow_pred_to_x0( x0 = xt - sigma_t * flow_pred return x0.to(original_dtype) + @staticmethod + def _build_i2v_model_input_writer( + latent_chunk: torch.Tensor, + condition_chunk: torch.Tensor, + target_dtype: torch.dtype, + ) -> Callable[[torch.Tensor], torch.Tensor]: + """Allocate one I2V input buffer and overwrite only its latent channels.""" + batch, latent_channels, frames, height, width = latent_chunk.shape + condition = condition_chunk.to(device=latent_chunk.device, dtype=target_dtype) + model_input = torch.empty( + (batch, latent_channels + condition.shape[1], frames, height, width), + dtype=target_dtype, + device=latent_chunk.device, + ) + model_input[:, latent_channels:].copy_(condition) + + def write(current_latent: torch.Tensor) -> torch.Tensor: + model_input[:, :latent_channels].copy_(current_latent) + return model_input + + return write + def denoise_chunk( self, latent_chunk: torch.Tensor, @@ -401,9 +502,18 @@ def denoise_chunk( current_start: int, max_attention_size: int, generator: torch.Generator | None = None, + session_input_cache: dict[str, object] | None = None, + prepared_control_is_sharded: bool = False, + prepare_model_input: Callable[[torch.Tensor], torch.Tensor] | None = None, + benchmark_events: list[tuple[torch.cuda.Event, torch.cuda.Event]] | None = None, ) -> torch.Tensor: current_latent = latent_chunk for timestep_idx in range(len(timesteps)): + step_start = step_end = None + if benchmark_events is not None: + step_start = torch.cuda.Event(enable_timing=True) + step_end = torch.cuda.Event(enable_timing=True) + step_start.record() schedule_timestep = timesteps[timestep_idx].view(1).to(device=current_latent.device) model_timestep = schedule_timestep.to(dtype=torch.float32) with torch.amp.autocast( @@ -411,17 +521,27 @@ def denoise_chunk( dtype=self.torch_dtype, enabled=current_latent.device.type == "cuda", ): + if prepare_model_input is None: + model_input = current_latent.to(dtype=self.torch_dtype) + model_condition = condition_chunk + else: + model_input = prepare_model_input(current_latent) + model_condition = None noise_pred = self.dit( - x=current_latent.to(dtype=self.torch_dtype), + x=model_input, timestep=model_timestep, context=prompt_emb, - y=condition_chunk, + y=model_condition, control_tensor=control_chunk, kv_cache=self_kv_cache, crossattn_cache=crossattn_cache, current_start=current_start, max_attention_size=max_attention_size, + _session_input_cache=session_input_cache, + _prepared_control_is_sharded=prepared_control_is_sharded or timestep_idx > 0, ) + if noise_pred is None: + raise RuntimeError("LingBot DMD forward unexpectedly returned no prediction") x0 = self._convert_flow_pred_to_x0(noise_pred, current_latent, schedule_timestep[0], scheduler) if timestep_idx < len(timesteps) - 1: next_timestep = timesteps[timestep_idx + 1].view(1).to(device=x0.device) @@ -429,6 +549,9 @@ def denoise_chunk( current_latent = scheduler.add_noise(x0, noise, next_timestep) else: current_latent = x0 + if benchmark_events is not None: + step_end.record() + benchmark_events.append((step_start, step_end)) logger.debug("LingBotWorldFast chunk denoised") return current_latent @@ -442,51 +565,136 @@ def _next_noise_chunk(self, state: _DenoisingCacheState) -> torch.Tensor: dtype=torch.float32, ) + @staticmethod + def _tensor_cache_key(tensor: torch.Tensor) -> tuple[object, ...]: + """Return mutation-sensitive identity facts without retaining a CPU copy.""" + return ( + id(tensor), + tuple(tensor.shape), + tensor.dtype, + tensor.device, + tensor._version, + ) + + def _prepare_session_inputs( + self, + state: _DenoisingCacheState, + cache_handle: int, + prompt_emb: torch.Tensor, + control_chunk: torch.Tensor | None, + ) -> None: + parameter = next(self.dit.parameters()) + text_key = (*self._tensor_cache_key(prompt_emb), id(parameter), parameter.dtype, parameter.device) + if state.projected_context_key != text_key: + state.projected_context_key = text_key + state.session_input_cache.pop("projected_context", None) + if control_chunk is None: + state.prepared_control_key = None + state.prepared_control_is_sharded = False + state.session_input_cache.pop("prepared_control", None) + return + key = (cache_handle, *self._tensor_cache_key(control_chunk)) + if state.prepared_control_key == key: + return + state.prepared_control_key = key + state.prepared_control_is_sharded = False + state.session_input_cache.pop("prepared_control", None) + @with_model_offload(["dit"]) def denoise_and_update_cache( self, cache_handle: int, condition_chunk: torch.Tensor, - prompt_emb: torch.Tensor, + prompt_emb: torch.Tensor | None, control_chunk: torch.Tensor | None, current_start: int, max_attention_size: int, - ) -> torch.Tensor: + _local_vae_handoff: bool = False, + _benchmark_profile: bool = False, + ) -> torch.Tensor | None | tuple[torch.Tensor | None, dict[str, object]]: """Denoise a chunk and commit its clean KV state inside each worker.""" try: state = self._cache_registry[cache_handle] except KeyError as exc: raise KeyError(f"Unknown cache handle {cache_handle}") from exc - denoised = self.denoise_chunk( - latent_chunk=self._next_noise_chunk(state), - condition_chunk=condition_chunk, - prompt_emb=prompt_emb, - timesteps=state.timesteps, - scheduler=state.scheduler, - control_chunk=control_chunk, - self_kv_cache=state.self_kv_cache, - crossattn_cache=state.crossattn_cache, - current_start=current_start, - max_attention_size=max_attention_size, - generator=state.generator, - ) - with torch.amp.autocast( - self.device.type, - dtype=self.torch_dtype, - enabled=self.device.type == "cuda", - ): - self.dit( - x=denoised.to(dtype=self.torch_dtype), - timestep=torch.zeros((1,), dtype=torch.float32, device=self.device), - context=prompt_emb, - y=condition_chunk, - control_tensor=control_chunk, - kv_cache=state.self_kv_cache, + session_prompt_emb = state.prompt_emb if state.prompt_emb is not None else prompt_emb + if session_prompt_emb is None: + raise ValueError("LingBot denoising requires a session prompt embedding") + self._prepare_session_inputs(state, cache_handle, session_prompt_emb, control_chunk) + try: + latent_chunk = self._next_noise_chunk(state) + prepare_model_input = self._build_i2v_model_input_writer( + latent_chunk, + condition_chunk, + self.torch_dtype, + ) + benchmark_events = [] if _benchmark_profile and self.device.type == "cuda" else None + denoise_span_start = None + if benchmark_events is not None: + denoise_span_start = torch.cuda.Event(enable_timing=True) + denoise_span_start.record() + denoised = self.denoise_chunk( + latent_chunk=latent_chunk, + condition_chunk=condition_chunk, + prompt_emb=session_prompt_emb, + timesteps=state.timesteps, + scheduler=state.scheduler, + control_chunk=control_chunk, + self_kv_cache=state.self_kv_cache, crossattn_cache=state.crossattn_cache, current_start=current_start, max_attention_size=max_attention_size, + generator=state.generator, + session_input_cache=state.session_input_cache, + prepared_control_is_sharded=state.prepared_control_is_sharded, + prepare_model_input=prepare_model_input, + benchmark_events=benchmark_events, ) - return denoised + state.prepared_control_is_sharded = control_chunk is not None + clean_start = clean_end = None + if benchmark_events is not None: + clean_start = torch.cuda.Event(enable_timing=True) + clean_end = torch.cuda.Event(enable_timing=True) + clean_start.record() + with torch.amp.autocast( + self.device.type, + dtype=self.torch_dtype, + enabled=self.device.type == "cuda", + ): + self.dit( + x=prepare_model_input(denoised), + timestep=torch.zeros((1,), dtype=torch.float32, device=self.device), + context=session_prompt_emb, + y=None, + control_tensor=control_chunk, + kv_cache=state.self_kv_cache, + crossattn_cache=state.crossattn_cache, + current_start=current_start, + max_attention_size=max_attention_size, + _session_input_cache=state.session_input_cache, + _prepared_control_is_sharded=state.prepared_control_is_sharded, + update_cache_only=True, + ) + profile = None + if benchmark_events is not None: + clean_end.record() + clean_end.synchronize() + profile = { + "dmd_step_seconds": [start.elapsed_time(end) / 1000.0 for start, end in benchmark_events], + "clean_kv_seconds": clean_start.elapsed_time(clean_end) / 1000.0, + "denoise_gpu_span_seconds": denoise_span_start.elapsed_time(clean_end) / 1000.0, + } + if _local_vae_handoff and self._vae_decode_stage is not None: + self._pending_vae_decode_latents.setdefault(cache_handle, deque()).append(denoised) + result = None + else: + result = denoised + return (result, profile) if profile is not None else result + finally: + # Camera tensors are immutable only within this chunk and can be large. + state.prepared_control_key = None + state.prepared_control_is_sharded = False + state.session_input_cache.pop("prepared_control", None) def advance_noise(self, cache_handle: int) -> bool: """Advance the actor-owned noise RNG for a decode-only cache hit.""" @@ -508,6 +716,9 @@ def list_cache_handles(self) -> tuple[int, ...]: def release_cache(self, cache_handle: int) -> bool: """Idempotently release worker-local state for one generation session.""" state = self._cache_registry.pop(cache_handle, None) + pending_latents = getattr(self, "_pending_vae_decode_latents", None) + if pending_latents is not None: + pending_latents.pop(cache_handle, None) if state is None: return False pool = getattr(self, "_cache_pool", None) diff --git a/tests/unit/models/test_lingbot_world_fast_dit.py b/tests/unit/models/test_lingbot_world_fast_dit.py index 4879f2d..f7cc5a1 100644 --- a/tests/unit/models/test_lingbot_world_fast_dit.py +++ b/tests/unit/models/test_lingbot_world_fast_dit.py @@ -51,6 +51,50 @@ def fake_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs: assert captured["output_layout"] == "BSND" +def test_causal_self_attention_packs_qkv_into_one_ulysses_collective() -> None: + attention = CausalSelfAttention(dim=32, num_heads=4) + freqs = precompute_freqs_cis_3d(8) + freqs_cos = torch.cat([freq.real for freq in freqs], dim=-1) + freqs_sin = torch.cat([freq.imag for freq in freqs], dim=-1) + cache = { + "k": torch.zeros(1, 12, 4, 8), + "v": torch.zeros(1, 12, 4, 8), + "global_end_index": 0, + "local_end_index": 0, + } + + def fake_scatter(tensor: torch.Tensor, _group: object): + return lambda: tensor + + def fake_gather(tensor: torch.Tensor, _group: object, *, num_heads: int): + assert num_heads == 4 + return lambda: tensor + + with ( + patch("telefuser.models.lingbot_world_fast_dit.get_ulysses_group", return_value=object()), + patch("telefuser.models.lingbot_world_fast_dit.get_ulysses_world_size", return_value=4), + patch( + "telefuser.models.lingbot_world_fast_dit.ulysses_scatter_heads", + side_effect=fake_scatter, + ) as scatter, + patch("telefuser.models.lingbot_world_fast_dit.ulysses_gather_heads", side_effect=fake_gather), + patch("telefuser.models.lingbot_world_fast_dit.attn_func", side_effect=lambda q, _k, _v, **_kwargs: q), + ): + attention( + torch.randn(1, 4, 32), + freqs_cos, + freqs_sin, + (1, 2, 2), + cache, + current_start=0, + max_attention_size=4, + ) + + scatter.assert_called_once() + packed_qkv = scatter.call_args.args[0] + assert packed_qkv.shape == (1, 4, 4, 24) + + def test_cached_cross_attention_uses_unified_attention_and_bsnd_cache() -> None: attention = CachedCrossAttention(dim=32, num_heads=4) attention_config = AttentionConfig.dense_attention(AttnImplType.SAGE_ATTN_2_8_8_SM90) @@ -115,3 +159,202 @@ def test_set_attention_config_updates_all_blocks() -> None: for block in model.blocks: assert block.self_attn.attention_config is attention_config assert block.cross_attn.attention_config is attention_config + + +def test_scalar_timestep_modulation_stays_broadcastable() -> None: + model = LingBotWorldFastDiT( + in_dim=4, + dim=32, + ffn_dim=64, + freq_dim=8, + text_dim=16, + out_dim=4, + num_heads=4, + num_layers=1, + ) + + scalar_head, scalar_modulation = model._build_timestep_embeddings(torch.tensor([5.0]), seq_len=3) + token_head, token_modulation = model._build_timestep_embeddings(torch.full((1, 3), 5.0), seq_len=3) + + assert scalar_head.shape == (1, 1, 32) + assert scalar_modulation.shape == (1, 6, 32) + assert token_head.shape == (1, 3, 32) + assert token_modulation.shape == (1, 3, 6, 32) + torch.testing.assert_close(token_head, scalar_head.expand_as(token_head)) + torch.testing.assert_close(token_modulation, scalar_modulation.unsqueeze(1).expand_as(token_modulation)) + + +def test_camera_conditioner_runs_after_sequence_sharding() -> None: + model = LingBotWorldFastDiT( + patch_size=(1, 2, 2), + in_dim=4, + dim=32, + ffn_dim=64, + freq_dim=8, + text_dim=16, + out_dim=4, + num_heads=4, + num_layers=2, + ).eval() + model.device_mesh = object() + control = torch.randn(1, 6 * 64, 1, 4, 4) + + def fake_shard(_mesh: object, tensors: list[torch.Tensor], dims: list[int]) -> None: + assert dims == [1] + local = tensors[0][:, :1].clone() + tensors[0].resize_(local.shape) + tensors[0].copy_(local) + + with ( + torch.inference_mode(), + patch( + "telefuser.models.lingbot_world_fast_dit.sequence_parallel_shard", + side_effect=fake_shard, + ) as shard, + ): + prepared = model._prepare_control(control, shard_for_usp=True) + + assert prepared is not None + control_tokens, camera_modulations = prepared + assert control_tokens.shape == (1, 1, 32) + assert len(camera_modulations) == 2 + assert all(scale.shape == shift.shape == (1, 1, 32) for scale, shift in camera_modulations) + shard.assert_called_once() + + +def _kv_cache() -> tuple[list[dict[str, torch.Tensor | int]], list[dict[str, torch.Tensor | bool | int]]]: + self_cache = [ + { + "k": torch.zeros(1, 1, 4, 8), + "v": torch.zeros(1, 1, 4, 8), + "global_end_index": 0, + "local_end_index": 0, + } + ] + cross_cache = [ + { + "k": torch.empty(1, 2, 4, 8), + "v": torch.empty(1, 2, 4, 8), + "is_init": False, + "sequence_length": 0, + } + ] + return self_cache, cross_cache + + +def test_prepared_control_and_cache_only_forward_preserve_transformer_execution() -> None: + model = LingBotWorldFastDiT( + patch_size=(1, 2, 2), + in_dim=8, + dim=32, + ffn_dim=64, + freq_dim=8, + text_dim=16, + out_dim=4, + num_heads=4, + num_layers=1, + ).eval() + x = torch.randn(1, 4, 1, 2, 2) + condition = torch.randn_like(x) + context = torch.randn(1, 2, 16) + control = torch.randn(1, 6 * 64, 1, 2, 2) + timestep = torch.zeros(1) + + self_cache, cross_cache = _kv_cache() + baseline = model( + x, + timestep, + context, + y=condition, + control_tensor=control, + kv_cache=self_cache, + crossattn_cache=cross_cache, + max_attention_size=1, + ) + prepared_control = model._prepare_control(control) + projected_context = model._project_text_context(context) + self_cache, cross_cache = _kv_cache() + cached = model( + x, + timestep, + context, + y=condition, + control_tensor=control, + kv_cache=self_cache, + crossattn_cache=cross_cache, + max_attention_size=1, + _prepared_control=prepared_control, + _projected_context=projected_context, + ) + + torch.testing.assert_close(cached, baseline, rtol=0, atol=0) + + session_input_cache: dict[str, object] = {} + with ( + patch.object(model, "_project_text_context", wraps=model._project_text_context) as project_context, + patch.object(model, "_prepare_control", wraps=model._prepare_control) as prepare_control, + patch.object( + model.blocks[0].self_attn, + "_prepare_causal_rope", + wraps=model.blocks[0].self_attn._prepare_causal_rope, + ) as prepare_causal_rope, + ): + self_cache, cross_cache = _kv_cache() + first_cached = model( + x, + timestep, + context, + y=condition, + control_tensor=control, + kv_cache=self_cache, + crossattn_cache=cross_cache, + max_attention_size=1, + _session_input_cache=session_input_cache, + ) + self_cache, cross_cache = _kv_cache() + second_cached = model( + x, + timestep, + context, + y=condition, + control_tensor=control, + kv_cache=self_cache, + crossattn_cache=cross_cache, + max_attention_size=1, + _session_input_cache=session_input_cache, + ) + + torch.testing.assert_close(first_cached, baseline, rtol=0, atol=0) + torch.testing.assert_close(second_cached, baseline, rtol=0, atol=0) + project_context.assert_called_once() + prepare_control.assert_called_once() + prepare_causal_rope.assert_called_once() + + self_cache, cross_cache = _kv_cache() + with ( + patch.object(model.blocks[0].self_attn.o, "forward", wraps=model.blocks[0].self_attn.o.forward) as self_out, + patch.object(model.blocks[0].cross_attn, "forward", wraps=model.blocks[0].cross_attn.forward) as cross_attn, + patch.object(model.blocks[0].ffn, "forward", wraps=model.blocks[0].ffn.forward) as ffn, + patch.object(model.head, "forward", wraps=model.head.forward) as head_forward, + ): + output = model( + x, + timestep, + context, + y=condition, + control_tensor=control, + kv_cache=self_cache, + crossattn_cache=cross_cache, + max_attention_size=1, + _prepared_control=prepared_control, + _projected_context=projected_context, + update_cache_only=True, + ) + + assert output is None + self_out.assert_not_called() + cross_attn.assert_not_called() + ffn.assert_not_called() + head_forward.assert_not_called() + assert self_cache[0]["global_end_index"] == 1 + assert cross_cache[0]["is_init"] is False diff --git a/tests/unit/ops/test_attention_backends.py b/tests/unit/ops/test_attention_backends.py index 5b2f686..0394033 100644 --- a/tests/unit/ops/test_attention_backends.py +++ b/tests/unit/ops/test_attention_backends.py @@ -1,7 +1,63 @@ from types import ModuleType -from unittest.mock import patch +from unittest.mock import MagicMock, patch -from telefuser.ops.attention import backends +import torch + +from telefuser.ops.attention import attention_impl, backends + + +def test_flash_attn4_dispatch_uses_cute_return_lse_argument() -> None: + q = torch.randn(1, 3, 2, 4) + flash_attn4 = MagicMock(return_value=(q, torch.zeros(1, 2, 3))) + + with ( + patch.object(attention_impl, "FLASH_ATTN_4_AVAILABLE", True), + patch.object(attention_impl, "flash_attn4", flash_attn4), + ): + output, lse = attention_impl.attention( + q, + q, + q, + attn_impl=attention_impl.AttnImplType.FLASH_ATTN_4, + return_lse=True, + ) + + assert output is q + assert lse.shape == (1, 2, 3) + flash_attn4.assert_called_once_with( + q, + q, + q, + softmax_scale=None, + causal=False, + return_lse=True, + ) + + +def test_flash_attn4_dispatch_unwraps_output_when_lse_is_disabled() -> None: + q = torch.randn(1, 3, 2, 4) + flash_attn4 = MagicMock(return_value=(q, None)) + + with ( + patch.object(attention_impl, "FLASH_ATTN_4_AVAILABLE", True), + patch.object(attention_impl, "flash_attn4", flash_attn4), + ): + output = attention_impl.attention( + q, + q, + q, + attn_impl=attention_impl.AttnImplType.FLASH_ATTN_4, + ) + + assert output is q + flash_attn4.assert_called_once_with( + q, + q, + q, + softmax_scale=None, + causal=False, + return_lse=False, + ) def test_sage_attention_prefers_tf_kernel() -> None: diff --git a/tests/unit/ops/test_normalization.py b/tests/unit/ops/test_normalization.py index 41a71d9..24abd92 100644 --- a/tests/unit/ops/test_normalization.py +++ b/tests/unit/ops/test_normalization.py @@ -4,7 +4,38 @@ import torch import torch.nn as nn -from telefuser.ops.normalization import AdaLayerNormContinuous, LayerNorm, RMSNorm +from telefuser.ops.normalization import ( + AdaLayerNormContinuous, + LayerNorm, + RMSNorm, + _fused_add_layer_norm_scale_shift, + _fused_layer_norm_scale_shift, +) + + +def test_fused_layer_norm_scale_shift_matches_native() -> None: + x = torch.randn(2, 5, 16) + scale = torch.randn(2, 1, 16) + shift = torch.randn(2, 1, 16) + + output = _fused_layer_norm_scale_shift(x, scale, shift) + expected = torch.nn.functional.layer_norm(x, (16,)) * (1 + scale) + shift + + torch.testing.assert_close(output, expected, atol=5e-5, rtol=1e-5) + + +def test_fused_add_layer_norm_scale_shift_matches_native() -> None: + residual = torch.randn(2, 5, 16) + x = torch.randn_like(residual) + scale = torch.randn(2, 1, 16) + shift = torch.randn(2, 1, 16) + + output, residual_out = _fused_add_layer_norm_scale_shift(residual, x, scale, shift) + expected_residual = residual + x + expected = torch.nn.functional.layer_norm(expected_residual, (16,)) * (1 + scale) + shift + + torch.testing.assert_close(residual_out, expected_residual) + torch.testing.assert_close(output, expected, atol=5e-5, rtol=1e-5) class TestRMSNorm: diff --git a/tests/unit/pipelines/lingbot_world_fast/test_parallelism.py b/tests/unit/pipelines/lingbot_world_fast/test_parallelism.py index 2b10d54..05d45ba 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_parallelism.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_parallelism.py @@ -26,6 +26,7 @@ def test_denoising_stage_parallel_models_enables_ulysses_and_fsdp() -> None: module_manager = MagicMock() module_manager.fetch_module.return_value = dit stage = LingBotWorldFastDenoisingStage("denoise", module_manager, runtime_config) + assert stage.empty_cache_after_call is False device_mesh = MagicMock() fsdp_model = MagicMock() diff --git a/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py b/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py index 8e2353f..bda0b4e 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py @@ -301,6 +301,25 @@ def denoise(active_generator: torch.Generator) -> torch.Tensor: assert not torch.equal(first, second) +def test_i2v_model_input_writer_reuses_storage_and_preserves_condition() -> None: + latent = torch.zeros(1, 2, 3, 2, 2, dtype=torch.float32) + condition = torch.arange(36, dtype=torch.float32).reshape(1, 3, 3, 2, 2) + write = LingBotWorldFastDenoisingStage._build_i2v_model_input_writer( + latent, + condition, + torch.bfloat16, + ) + + first = write(torch.ones_like(latent)) + storage_pointer = first.data_ptr() + second = write(torch.full_like(latent, 2.0)) + + assert second.data_ptr() == storage_pointer + assert second.dtype == torch.bfloat16 + torch.testing.assert_close(second[:, :2].float(), torch.full_like(latent, 2.0)) + torch.testing.assert_close(second[:, 2:].float(), condition) + + def test_runtime_truncates_non_aligned_latent_frame_count() -> None: _, runtime = _create_runtime(frame_num=13) diff --git a/tests/unit/pipelines/lingbot_world_fast/test_session_cache.py b/tests/unit/pipelines/lingbot_world_fast/test_session_cache.py index efe8b3a..38e0f99 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_session_cache.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_session_cache.py @@ -1,3 +1,4 @@ +from collections import deque from types import SimpleNamespace from unittest.mock import MagicMock @@ -18,7 +19,11 @@ def _cache_stage() -> LingBotWorldFastDenoisingStage: return stage -def _initialize_cache(stage: LingBotWorldFastDenoisingStage, cache_handle: int) -> bool: +def _initialize_cache( + stage: LingBotWorldFastDenoisingStage, + cache_handle: int, + prompt_emb: torch.Tensor | None = None, +) -> bool: generator_state = torch.Generator(device="cpu").manual_seed(cache_handle).get_state().tolist() noise_generator_state = torch.Generator(device="cpu").manual_seed(cache_handle + 100).get_state().tolist() return LingBotWorldFastDenoisingStage.initialize_cache.__wrapped__( @@ -31,9 +36,19 @@ def _initialize_cache(stage: LingBotWorldFastDenoisingStage, cache_handle: int) generator_state=generator_state, noise_generator_state=noise_generator_state, noise_shape=(1, 16, 1, 1, 1), + prompt_emb=prompt_emb, ) +def test_worker_cache_retains_session_prompt_embedding() -> None: + stage = _cache_stage() + prompt_emb = torch.randn(1, 8, 16) + + assert _initialize_cache(stage, 11, prompt_emb=prompt_emb) is True + + assert stage._cache_registry[11].prompt_emb is prompt_emb + + def test_worker_cache_registry_isolates_handles_and_releases_idempotently() -> None: stage = _cache_stage() @@ -123,6 +138,65 @@ def test_worker_owned_noise_rng_advances_deterministically() -> None: torch.testing.assert_close(replicated_third, expected_third) +class _SessionCacheTestDiT(nn.Module): + def __init__(self) -> None: + super().__init__() + self.text_embedding = nn.Sequential(nn.Linear(2, 2)) + + +def test_session_input_caches_invalidate_on_prompt_and_control_mutation() -> None: + stage = LingBotWorldFastDenoisingStage.__new__(LingBotWorldFastDenoisingStage) + stage.dit = _SessionCacheTestDiT() + state = SimpleNamespace( + projected_context_key=None, + prepared_control_key=None, + prepared_control_is_sharded=False, + session_input_cache={}, + ) + prompt = torch.zeros(1, 2) + control = torch.zeros(1, 2) + + stage._prepare_session_inputs(state, 9, prompt, control) + cached_context = torch.ones(1, 2) + cached_control = (torch.ones(1, 2), ()) + state.session_input_cache.update(projected_context=cached_context, prepared_control=cached_control) + stage._prepare_session_inputs(state, 9, prompt, control) + + assert state.session_input_cache["projected_context"] is cached_context + assert state.session_input_cache["prepared_control"] is cached_control + + prompt.add_(1) + stage._prepare_session_inputs(state, 9, prompt, control) + + assert "projected_context" not in state.session_input_cache + assert state.session_input_cache["prepared_control"] is cached_control + + control.add_(1) + stage._prepare_session_inputs(state, 9, prompt, control) + + assert "prepared_control" not in state.session_input_cache + + +def test_colocated_decode_consumes_worker_local_latent() -> None: + stage = LingBotWorldFastDenoisingStage.__new__(LingBotWorldFastDenoisingStage) + decoder = MagicMock() + decoder.decode_chunk.return_value = torch.ones(1) + stage._vae_decode_stage = decoder + local_latent = torch.ones(1, 2, 3) + next_local_latent = torch.full((1, 2, 3), 2.0) + stage._pending_vae_decode_latents = {9: deque([local_latent, next_local_latent])} + + output = stage.decode_chunk(9, torch.empty(0, dtype=torch.uint8), True, False) + stage.decode_chunk(9, torch.empty(0, dtype=torch.uint8), False, True) + + assert torch.equal(output, torch.ones(1)) + assert decoder.decode_chunk.call_args_list == [ + ((9, local_latent, True, False),), + ((9, next_local_latent, False, True),), + ] + assert stage._pending_vae_decode_latents == {} + + class _RecordingDecoder(nn.Module): def __init__(self) -> None: super().__init__() From 25f2e73dc1d297183cc3dbd334d14604048c8e3e Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Sun, 2 Aug 2026 06:27:03 +0000 Subject: [PATCH 03/11] perf(lingbot): colocate spatial VAE decode Shard Wan VAE decode over height with halo exchange, colocate four-GPU decode with the DiT workers, and reuse local latent and shared CPU output buffers. Serialize overlapping DiT/decode collectives, streamline worker queues and thread pools, preserve synchronized phase profiles, and keep the documented four- and six-GPU placements covered by tests. --- docs/en/stream_scheduler.md | 8 +- .../lingbot_world_fast_image_to_video_h100.py | 25 ++- .../lingbot_world_v2_image_to_video_h100.py | 60 +++++-- telefuser/distributed/vae_spatial.py | 169 ++++++++++++++++++ telefuser/models/wan_video_vae.py | 104 ++++++++++- .../pipelines/lingbot_world_fast/pipeline.py | 88 ++++++++- .../pipelines/lingbot_world_fast/service.py | 18 +- .../pipelines/lingbot_world_fast/streaming.py | 116 ++++++++++-- .../pipelines/lingbot_world_fast/vae_stage.py | 123 +++++++++---- telefuser/worker/parallel_worker.py | 28 +-- .../unit/models/test_wan_video_vae_spatial.py | 62 +++++++ .../test_pipeline_lifecycle.py | 64 ++++++- .../test_service_action_loop.py | 15 ++ .../lingbot_world_fast/test_stream_example.py | 35 +++- .../lingbot_world_fast/test_streaming.py | 13 ++ .../test_vae_stage_capacity.py | 70 +++++++- .../lingbot_world_v2/test_service.py | 48 +++-- tests/unit/worker/test_parallel_worker.py | 21 ++- 18 files changed, 952 insertions(+), 115 deletions(-) create mode 100644 telefuser/distributed/vae_spatial.py create mode 100644 tests/unit/models/test_wan_video_vae_spatial.py diff --git a/docs/en/stream_scheduler.md b/docs/en/stream_scheduler.md index 28e732e..60282c8 100644 --- a/docs/en/stream_scheduler.md +++ b/docs/en/stream_scheduler.md @@ -116,9 +116,11 @@ never transfer actor-owned stage state between workers. `StreamingResourceGroupSpec` represents an explicit shared concurrency constraint. A stage participates only when its `StreamingStageSpec.resource_group` names a group declared by `StreamingPipelineSpec.resource_groups`. -Do not infer a resource group from `device_id` or `ParallelConfig.device_ids`. For LingBot, VAE encode, DiT, and VAE -decode are independent actors and may overlap on the same GPU. If a placement exceeds memory capacity, move stages to -different devices or define a deliberate deployment constraint; do not add an implicit global mutex. +Do not infer a resource group from `device_id` or `ParallelConfig.device_ids`. LingBot VAE encode remains an +independent actor. When distributed DiT and VAE decode use exactly the same device list and world size, the pipeline +explicitly co-locates the decoder in the DiT worker group to reuse CUDA contexts; non-matching placements remain +independent actors. If a placement exceeds memory capacity, move stages to different devices or define a deliberate +deployment constraint; do not add an implicit global mutex. LingBot uses independent `vae_encode_config` and `vae_decode_config`. Each VAE stage receives its own complete `ModelRuntimeConfig`; there is no shared VAE diff --git a/examples/lingbot/lingbot_world_fast_image_to_video_h100.py b/examples/lingbot/lingbot_world_fast_image_to_video_h100.py index 10e356b..eb3d07d 100644 --- a/examples/lingbot/lingbot_world_fast_image_to_video_h100.py +++ b/examples/lingbot/lingbot_world_fast_image_to_video_h100.py @@ -5,6 +5,8 @@ Four GPUs with Ulysses sequence parallelism: python examples/lingbot/lingbot_world_fast_image_to_video_h100.py --gpu_num 4 +Six GPUs with five-way DiT parallelism and one dedicated VAE GPU: + python examples/lingbot/lingbot_world_fast_image_to_video_h100.py --gpu_num 6 LiveKit streaming service: telefuser stream-serve examples/lingbot/lingbot_world_fast_image_to_video_h100.py \ --livekit-url ws://127.0.0.1:7880 \ @@ -78,20 +80,22 @@ ) -def _resolve_stage_devices(total_gpu_count: int) -> tuple[list[int], int, int]: +def _resolve_stage_devices(total_gpu_count: int) -> tuple[list[int], int, list[int]]: """Return DiT, VAE encode, and VAE decode devices for available GPUs.""" if total_gpu_count < 1: raise ValueError(f"parallelism must be positive, got {total_gpu_count}") - if total_gpu_count in {2, 4}: - return list(range(total_gpu_count)), 0, 1 + if total_gpu_count == 2: + return [0, 1], 0, [1] + if total_gpu_count == 4: + return [0, 1, 2, 3], 0, [0, 1, 2, 3] if total_gpu_count == 5: - return [0, 1, 2, 3], 4, 4 + return [0, 1, 2, 3], 4, [4] if total_gpu_count == 6: - return [0, 1, 2, 3, 4], 5, 5 + return [0, 1, 2, 3, 4], 5, [5] return ( list(range(total_gpu_count)), int(PPL_CONFIG["vae_encode_device_id"]), - int(PPL_CONFIG["vae_decode_device_id"]), + [int(PPL_CONFIG["vae_decode_device_id"])], ) @@ -101,7 +105,7 @@ def get_pipeline( fast_model_root: str | None = None, ) -> LingBotWorldFastPipeline: """Load LingBot-World-Fast for offline chunked generation.""" - dit_device_ids, vae_encode_device, vae_decode_device = _resolve_stage_devices(parallelism) + dit_device_ids, vae_encode_device, vae_decode_device_ids = _resolve_stage_devices(parallelism) model_root_path = Path(model_root).expanduser() if model_root else None fast_model_root_path = Path(fast_model_root).expanduser() if fast_model_root else None vae_path = str(model_root_path / "Wan2.1_VAE.pth") if model_root_path else PPL_CONFIG["vae_path"] @@ -148,9 +152,12 @@ def get_pipeline( ), vae_decode_config=ModelRuntimeConfig( device_type="cuda", - device_id=vae_decode_device, + device_id=vae_decode_device_ids[0], torch_dtype=PPL_CONFIG["vae_torch_dtype"], - parallel_config=ParallelConfig(device_ids=[vae_decode_device]), + parallel_config=ParallelConfig( + device_ids=vae_decode_device_ids, + sp_ulysses_degree=len(vae_decode_device_ids), + ), ), text_encoding_config=ModelRuntimeConfig(device_type="cuda", device_id=dit_device_ids[0], torch_dtype=dtype), dit_config=ModelRuntimeConfig( diff --git a/examples/lingbot/lingbot_world_v2_image_to_video_h100.py b/examples/lingbot/lingbot_world_v2_image_to_video_h100.py index b28a56f..bb1d368 100644 --- a/examples/lingbot/lingbot_world_v2_image_to_video_h100.py +++ b/examples/lingbot/lingbot_world_v2_image_to_video_h100.py @@ -5,6 +5,8 @@ Four GPUs with Ulysses sequence parallelism: python examples/lingbot/lingbot_world_v2_image_to_video_h100.py --gpu_num 4 +Six GPUs with five-way DiT parallelism and one dedicated VAE GPU: + python examples/lingbot/lingbot_world_v2_image_to_video_h100.py --gpu_num 6 Multi-GPU runs configure the VAE worker and DiT SP group independently in PPL_CONFIG. LiveKit streaming service: telefuser stream-serve examples/lingbot/lingbot_world_v2_image_to_video_h100.py \ @@ -29,6 +31,7 @@ from telefuser.models.lingbot_world_fast_dit import LingBotWorldFastDiT from telefuser.models.wan_video_text_encoder import WanTextEncoder from telefuser.models.wan_video_vae import WanVideoVAE +from telefuser.ops.attention.backends import FLASH_ATTN_3_AVAILABLE, FLASH_ATTN_4_AVAILABLE from telefuser.pipelines.lingbot_world_fast.service import LingBotWorldFastService from telefuser.pipelines.lingbot_world_fast.session import LingBotWorldFastSessionConfig, resolve_lingbot_frame_count from telefuser.pipelines.lingbot_world_v2 import ( @@ -49,7 +52,8 @@ "A serene lakeside scene with a lone tree standing in calm water, surrounded by distant snow-capped " "mountains under a bright blue sky with drifting white clouds. Gentle ripples reflect the tree and sky." ) -RESOLUTION_AREAS = {"480p": 480 * 832, "720p": 720 * 1280} +RESOLUTION_SIZES = {"480p": (832, 480), "720p": (1280, 720)} +RESOLUTION_AREAS = {name: width * height for name, (width, height) in RESOLUTION_SIZES.items()} PPL_CONFIG = dict( vae_path=str(TF_MODEL_ZOO_PATH / "Wan2.2-I2V-A14B" / "Wan2.1_VAE.pth"), @@ -76,8 +80,14 @@ seed=42, target_fps=16, max_duration_seconds=120.0, - attn_impl=AttnImplType.SAGE_ATTN_2_8_8_SM90, - compile_config=CompileConfig(enabled=True), + attn_impl=( + AttnImplType.FLASH_ATTN_4 + if FLASH_ATTN_4_AVAILABLE + else AttnImplType.FLASH_ATTN_3 + if FLASH_ATTN_3_AVAILABLE + else AttnImplType.SAGE_ATTN_2_8_8_SM90 + ), + compile_config=CompileConfig(enabled=False), enable_fsdp=False, local_attn_size=18, sink_size=6, @@ -93,20 +103,34 @@ ) -def _resolve_stage_devices(total_gpu_count: int) -> tuple[list[int], int, int]: +class _LingBotWorldV2Service(LingBotWorldFastService): + """Normalize service images to the checkpoint configured output size.""" + + @staticmethod + def _load_image(config: dict) -> Image.Image: + image = LingBotWorldFastService._load_image(config) + target_size = RESOLUTION_SIZES[PPL_CONFIG["resolution"]] + if image.size != target_size: + image = image.resize(target_size, Image.Resampling.BICUBIC) + return image + + +def _resolve_stage_devices(total_gpu_count: int) -> tuple[list[int], int, list[int]]: """Return DiT, VAE encode, and VAE decode devices for available GPUs.""" if total_gpu_count < 1: raise ValueError(f"parallelism must be positive, got {total_gpu_count}") - if total_gpu_count in {2, 4}: - return list(range(total_gpu_count)), 0, 1 + if total_gpu_count == 2: + return [0, 1], 0, [1] + if total_gpu_count == 4: + return [0, 1, 2, 3], 0, [0, 1, 2, 3] if total_gpu_count == 5: - return [0, 1, 2, 3], 4, 4 + return [0, 1, 2, 3], 4, [4] if total_gpu_count == 6: - return [0, 1, 2, 3, 4], 5, 5 + return [0, 1, 2, 3, 4], 5, [5] return ( list(range(total_gpu_count)), int(PPL_CONFIG["vae_encode_device_id"]), - int(PPL_CONFIG["vae_decode_device_id"]), + [int(PPL_CONFIG["vae_decode_device_id"])], ) @@ -116,7 +140,7 @@ def get_pipeline( v2_model_root: str | None = None, ) -> LingBotWorldV2Pipeline: """Load LingBot-World v2 for offline chunked generation.""" - dit_device_ids, vae_encode_device, vae_decode_device = _resolve_stage_devices(parallelism) + dit_device_ids, vae_encode_device, vae_decode_device_ids = _resolve_stage_devices(parallelism) model_root_path = Path(model_root).expanduser() if model_root else None v2_model_root_path = Path(v2_model_root).expanduser() if v2_model_root else None @@ -157,6 +181,7 @@ def get_pipeline( pipeline.init( module_manager, LingBotWorldV2PipelineConfig( + max_area=RESOLUTION_AREAS[PPL_CONFIG["resolution"]] + 1, vae_encode_config=ModelRuntimeConfig( device_type="cuda", device_id=vae_encode_device, @@ -165,9 +190,12 @@ def get_pipeline( ), vae_decode_config=ModelRuntimeConfig( device_type="cuda", - device_id=vae_decode_device, + device_id=vae_decode_device_ids[0], torch_dtype=PPL_CONFIG["vae_torch_dtype"], - parallel_config=ParallelConfig(device_ids=[vae_decode_device]), + parallel_config=ParallelConfig( + device_ids=vae_decode_device_ids, + sp_ulysses_degree=len(vae_decode_device_ids), + ), ), text_encoding_config=ModelRuntimeConfig(device_type="cuda", device_id=dit_device_ids[0], torch_dtype=dtype), dit_config=ModelRuntimeConfig( @@ -190,7 +218,7 @@ def get_pipeline( def get_service(gpu_num: int = PPL_CONFIG["parallelism"]) -> LingBotWorldFastService: """Build the service loaded by the TeleFuser stream server.""" pipeline = get_pipeline(parallelism=gpu_num) - return LingBotWorldFastService( + return _LingBotWorldV2Service( pipeline, default_fps=PPL_CONFIG["target_fps"], max_generation_seconds=PPL_CONFIG["max_duration_seconds"], @@ -229,7 +257,10 @@ def run( """Generate a complete offline video through the pipeline core API.""" if resolution not in RESOLUTION_AREAS: raise ValueError(f"Unsupported resolution: {resolution}") - pipeline.config.max_area = RESOLUTION_AREAS[resolution] + pipeline.config.max_area = RESOLUTION_AREAS[resolution] + 1 + target_size = RESOLUTION_SIZES[resolution] + if image.size != target_size: + image = image.resize(target_size, Image.Resampling.BICUBIC) fps = PPL_CONFIG["target_fps"] if fps is None else fps session_config = LingBotWorldFastSessionConfig( @@ -318,7 +349,6 @@ def main( image = Image.open(image_path).convert("RGB") start = time.perf_counter() - frames = run( pipeline, image, diff --git a/telefuser/distributed/vae_spatial.py b/telefuser/distributed/vae_spatial.py new file mode 100644 index 0000000..2c4e2b3 --- /dev/null +++ b/telefuser/distributed/vae_spatial.py @@ -0,0 +1,169 @@ +"""Internal height-sharded helpers for causal VAE decoding.""" + +from __future__ import annotations + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F + + +def _spatial_rank() -> int: + return dist.get_rank() if dist.is_initialized() else 0 + + +def _spatial_world_size() -> int: + return dist.get_world_size() if dist.is_initialized() else 1 + + +def _height_memory_format(tensor: torch.Tensor) -> torch.memory_format: + if tensor.dim() == 5 and tensor.stride(1) == 1: + return torch.channels_last_3d + if tensor.dim() == 4 and tensor.stride(1) == 1: + return torch.channels_last + return torch.contiguous_format + + +def _split_height(tensor: torch.Tensor) -> torch.Tensor: + world_size = _spatial_world_size() + if world_size == 1: + return tensor + shard = torch.tensor_split(tensor, world_size, dim=-2)[_spatial_rank()] + return shard.contiguous(memory_format=_height_memory_format(tensor)) + + +def _gather_height_sizes(tensor: torch.Tensor) -> list[int]: + world_size = _spatial_world_size() + if world_size == 1: + return [tensor.shape[-2]] + local_height = torch.tensor([tensor.shape[-2]], dtype=torch.int64, device=tensor.device) + gathered = [torch.empty_like(local_height) for _ in range(world_size)] + dist.all_gather(gathered, local_height) + return [int(height.item()) for height in gathered] + + +def _gather_height(tensor: torch.Tensor) -> torch.Tensor: + world_size = _spatial_world_size() + if world_size == 1: + return tensor + heights = _gather_height_sizes(tensor) + max_height = max(heights) + if tensor.shape[-2] < max_height: + shape = list(tensor.shape) + shape[-2] = max_height - tensor.shape[-2] + tensor = torch.cat([tensor, tensor.new_zeros(shape)], dim=-2) + tensor = tensor.contiguous() + gathered = [torch.empty_like(tensor) for _ in range(world_size)] + dist.all_gather(gathered, tensor) + return torch.cat([shard[..., :height, :] for shard, height in zip(gathered, heights)], dim=-2) + + +def _ensure_halo_buffer(buffer: torch.Tensor | None, reference: torch.Tensor) -> torch.Tensor: + memory_format = _height_memory_format(reference) + if ( + buffer is None + or buffer.shape != reference.shape + or buffer.dtype != reference.dtype + or buffer.device != reference.device + or not buffer.is_contiguous(memory_format=memory_format) + ): + return torch.empty(reference.shape, dtype=reference.dtype, device=reference.device, memory_format=memory_format) + return buffer + + +def _exchange_height_halo(module: nn.Module, tensor: torch.Tensor, halo_size: int) -> torch.Tensor: + world_size = _spatial_world_size() + if world_size == 1 or halo_size == 0: + return tensor + rank = _spatial_rank() + top = tensor[..., :halo_size, :] + bottom = tensor[..., -halo_size:, :] + module._halo_recv_top = _ensure_halo_buffer(module._halo_recv_top, top) + module._halo_recv_bottom = _ensure_halo_buffer(module._halo_recv_bottom, bottom) + + operations = [] + if rank > 0: + operations.extend( + [ + dist.P2POp(dist.irecv, module._halo_recv_top, rank - 1), + dist.P2POp(dist.isend, top.contiguous(memory_format=_height_memory_format(top)), rank - 1), + ] + ) + if rank < world_size - 1: + operations.extend( + [ + dist.P2POp(dist.isend, bottom.contiguous(memory_format=_height_memory_format(bottom)), rank + 1), + dist.P2POp(dist.irecv, module._halo_recv_bottom, rank + 1), + ] + ) + for request in dist.batch_isend_irecv(operations): + request.wait() + + if rank == 0: + module._halo_recv_top.zero_() + if rank == world_size - 1: + module._halo_recv_bottom.zero_() + return torch.cat([module._halo_recv_top, tensor, module._halo_recv_bottom], dim=-2) + + +def _spatial_causal_conv3d_forward( + module: nn.Conv3d, + tensor: torch.Tensor, + cache_tensor: torch.Tensor | None, +) -> torch.Tensor: + padding = list(module._spatial_padding) + if cache_tensor is not None and padding[4] > 0: + cache_tensor = cache_tensor.to(tensor.device) + tensor = torch.cat([cache_tensor, tensor], dim=2) + padding[4] -= cache_tensor.shape[2] + if any(padding): + tensor = F.pad(tensor, padding) + tensor = _exchange_height_halo(module, tensor, module._height_halo_size) + tensor = tensor.contiguous(memory_format=torch.channels_last_3d) + return F.conv3d( + tensor, + module.weight, + module.bias, + module.stride, + (0, 0, 0), + module.dilation, + module.groups, + ) + + +class _SpatialParallelConv2d(nn.Conv2d): + """Conv2d over a height shard with one halo exchange per invocation.""" + + def __init__(self, source: nn.Conv2d) -> None: + if source.stride[0] != 1: + raise ValueError("VAE spatial decode only supports stride-one Conv2d layers") + if source.padding_mode != "zeros": + raise ValueError("VAE spatial decode only supports zero-padded Conv2d layers") + super().__init__( + source.in_channels, + source.out_channels, + source.kernel_size, + stride=source.stride, + padding=0, + dilation=source.dilation, + groups=source.groups, + bias=source.bias is not None, + padding_mode=source.padding_mode, + device=source.weight.device, + dtype=source.weight.dtype, + ) + self.weight = source.weight + self.bias = source.bias + self._height_halo_size = source.dilation[0] * (source.kernel_size[0] - 1) // 2 + if source.padding[0] != self._height_halo_size: + raise ValueError("VAE spatial Conv2d requires symmetric height padding") + self._width_padding = source.padding[1] + self._halo_recv_top: torch.Tensor | None = None + self._halo_recv_bottom: torch.Tensor | None = None + self.train(source.training) + + def forward(self, tensor: torch.Tensor) -> torch.Tensor: + if self._width_padding: + tensor = F.pad(tensor, (self._width_padding, self._width_padding, 0, 0)) + tensor = _exchange_height_halo(self, tensor, self._height_halo_size) + return F.conv2d(tensor, self.weight, self.bias, self.stride, 0, self.dilation, self.groups) diff --git a/telefuser/models/wan_video_vae.py b/telefuser/models/wan_video_vae.py index 915d081..6ffa7f2 100644 --- a/telefuser/models/wan_video_vae.py +++ b/telefuser/models/wan_video_vae.py @@ -10,6 +10,13 @@ from tqdm import tqdm from telefuser.core.base_model import BaseModel +from telefuser.distributed.vae_spatial import ( + _SpatialParallelConv2d, + _gather_height, + _spatial_causal_conv3d_forward, + _spatial_world_size, + _split_height, +) from telefuser.utils.logging import logger from telefuser.utils.model_weight import hash_state_dict_keys @@ -115,6 +122,49 @@ def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch return super().forward(x) +class _SpatialParallelCausalConv3d(CausalConv3d): + """Causal Conv3d over a height shard with neighboring halo exchange.""" + + def __init__(self, source: CausalConv3d) -> None: + temporal_padding = source._padding[4] // 2 + height_padding = source._padding[2] + width_padding = source._padding[0] + if source.stride[1] != 1: + raise ValueError("VAE spatial decode only supports stride-one height convolutions") + super().__init__( + source.in_channels, + source.out_channels, + source.kernel_size, + stride=source.stride, + padding=(temporal_padding, height_padding, width_padding), + dilation=source.dilation, + groups=source.groups, + bias=source.bias is not None, + padding_mode=source.padding_mode, + device=source.weight.device, + dtype=source.weight.dtype, + ) + self.weight = source.weight + self.bias = source.bias + self._height_halo_size = source.dilation[1] * (source.kernel_size[1] - 1) // 2 + if height_padding != self._height_halo_size: + raise ValueError("VAE spatial Conv3d requires symmetric height padding") + self._spatial_padding = ( + width_padding, + width_padding, + 0, + 0, + 2 * temporal_padding, + 0, + ) + self._halo_recv_top: torch.Tensor | None = None + self._halo_recv_bottom: torch.Tensor | None = None + self.train(source.training) + + def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor: + return _spatial_causal_conv3d_forward(self, x, cache_x) + + class RMS_norm(nn.Module): """RMS normalization with channel-first/last support.""" @@ -268,10 +318,13 @@ def __init__(self, dim: int): self.norm = RMS_norm(dim) self.to_qkv = nn.Conv2d(dim, dim * 3, 1) self.proj = nn.Conv2d(dim, dim, 1) + self._spatial_parallel = False nn.init.zeros_(self.proj.weight) def forward(self, x: torch.Tensor) -> torch.Tensor: + if self._spatial_parallel: + x = _gather_height(x).contiguous() identity = x b, c, t, h, w = x.size() x = rearrange(x, "b c t h w -> (b t) c h w") @@ -283,7 +336,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.proj(x) x = rearrange(x, "(b t) c h w-> b c t h w", t=t) - return x + identity + x = x + identity + return _split_height(x) if self._spatial_parallel else x class Encoder3d(nn.Module): @@ -414,6 +468,8 @@ def __init__( self.num_res_blocks = num_res_blocks self.attn_scales = attn_scales self.temperal_upsample = temperal_upsample + self._spatial_parallel = False + self._spatial_upsample_count = 0 dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] dims = [int(d * (1 - pruning_rate)) for d in dims] @@ -452,6 +508,10 @@ def __init__( def forward(self, x: torch.Tensor, feat_cache: list | None = None, feat_idx: list | None = None) -> torch.Tensor: """Forward pass with list-based feature caching.""" + expected_height = None + if self._spatial_parallel: + expected_height = x.shape[-2] * (2**self._spatial_upsample_count) + x = _split_height(x) if feat_cache is not None and feat_idx is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() @@ -498,9 +558,47 @@ def forward(self, x: torch.Tensor, feat_cache: list | None = None, feat_idx: lis feat_idx[0] += 1 else: x = layer(x) + if self._spatial_parallel: + x = _gather_height(x) + x = x[..., :expected_height, :].contiguous() return x +def _enable_spatial_parallel_decode(vae: WanVideoVAE) -> int: + """Replace decoder convolutions with height-sharded equivalents in a worker process.""" + if _spatial_world_size() <= 1: + return 0 + if getattr(vae, "parallelism", 1) > 1: + raise RuntimeError("Wan VAE native decode_parallel and streaming spatial decode are mutually exclusive") + decoder = vae.model.decoder + decoder._spatial_parallel = True + decoder._spatial_upsample_count = sum( + isinstance(module, Resample) and module.mode in {"upsample2d", "upsample3d"} for module in decoder.modules() + ) + converted = 0 + + def convert(module: nn.Module, inside_resample: bool = False) -> None: + nonlocal converted + if isinstance(module, AttentionBlock): + module._spatial_parallel = True + for name, child in list(module.named_children()): + if isinstance(child, (_SpatialParallelCausalConv3d, _SpatialParallelConv2d)): + continue + if isinstance(child, CausalConv3d): + setattr(module, name, _SpatialParallelCausalConv3d(child)) + converted += 1 + continue + child_inside_resample = inside_resample or isinstance(module, Resample) + if child_inside_resample and isinstance(child, nn.Conv2d): + setattr(module, name, _SpatialParallelConv2d(child)) + converted += 1 + continue + convert(child, child_inside_resample) + + convert(decoder) + return converted + + class VideoVAE(nn.Module): """Full VAE with encoder and decoder for video.""" @@ -742,7 +840,9 @@ def enable_channels_last_3d(self) -> int: """ return _convert_conv3d_to_channels_last_3d(self.model) - def set_parallelism(self, parallelism: int): + def set_parallelism(self, parallelism: int) -> None: + if parallelism > 1 and getattr(self.model.decoder, "_spatial_parallel", False): + raise RuntimeError("Wan VAE native decode_parallel and streaming spatial decode are mutually exclusive") self.parallelism = parallelism # ==================== 2D Spatial Parallel Methods ==================== diff --git a/telefuser/pipelines/lingbot_world_fast/pipeline.py b/telefuser/pipelines/lingbot_world_fast/pipeline.py index 149f101..f49b324 100644 --- a/telefuser/pipelines/lingbot_world_fast/pipeline.py +++ b/telefuser/pipelines/lingbot_world_fast/pipeline.py @@ -55,6 +55,66 @@ def _default_vae_stage_runtime_config() -> ModelRuntimeConfig: return ModelRuntimeConfig(torch_dtype=torch.float32, parallel_config=ParallelConfig(device_ids=[0])) +class _CoLocatedVAEDecodeWorker: + """Map VAE decode calls onto a denoising ParallelWorker.""" + + uses_local_latent_handoff = True + + def __init__(self, worker: ParallelWorker) -> None: + self._worker = worker + self.name = f"Co-located VAE decode on {worker.name}" + self._output_buffers: dict[int, torch.Tensor] = {} + + def reset_device_memory_peak(self) -> bool: + return self._worker.reset_vae_decode_device_memory_peak(sync=True) + + def device_memory_snapshots(self) -> list[dict[str, int | str]]: + return self._worker.vae_decode_device_memory_snapshots(sync=True) + + def estimate_session_cache_bytes(self) -> int: + return self._worker.estimate_vae_decode_session_cache_bytes(sync=True) + + def observed_session_cache_bytes(self) -> int: + return self._worker.observed_vae_decode_session_cache_bytes(sync=True) + + def configure_cache_pool(self, capacity: int) -> object: + return self._worker.configure_vae_decode_cache_pool(capacity, sync=True) + + def initialize_cache(self, cache_handle: int, *, sync: bool = False) -> object: + self._output_buffers.pop(cache_handle, None) + return self._worker.initialize_vae_decode_cache(cache_handle, sync=sync) + + def decode_chunk(self, cache_handle: int, *args: object, **kwargs: object) -> object: + result = self._worker.decode_chunk(cache_handle, *args, **kwargs) + + def resolve(value: object) -> object: + output, profile = value if isinstance(value, tuple) else (value, None) + if output is None or isinstance(output, torch.Tensor) and not output.numel(): + output = self._output_buffers.get(cache_handle) + if output is None: + raise RuntimeError(f"Co-located VAE output buffer {cache_handle} was not registered") + elif isinstance(output, torch.Tensor): + self._output_buffers[cache_handle] = output + else: + raise TypeError(f"Co-located VAE decode returned {type(output).__name__}, expected Tensor") + return (output, profile) if isinstance(value, tuple) else output + + return (lambda: resolve(result())) if callable(result) else resolve(result) + + def release_cache(self, cache_handle: int, *, sync: bool = False) -> object: + result = self._worker.release_vae_decode_cache(cache_handle, sync=sync) + if callable(result): + + def wait() -> object: + released = result() + self._output_buffers.pop(cache_handle, None) + return released + + return wait + self._output_buffers.pop(cache_handle, None) + return result + + @dataclass class LingBotWorldFastPipelineConfig: vae_encode_config: ModelRuntimeConfig = field(default_factory=_default_vae_stage_runtime_config) @@ -256,8 +316,8 @@ def init(self, module_manager: ModuleManager, config: LingBotWorldFastPipelineCo self.tokenizer = HuggingfaceTokenizer(tokenizer_path, 512, "whitespace") vae_encode_config = config.vae_encode_config vae_decode_config = config.vae_decode_config - self._validate_vae_stage_runtime_config(vae_encode_config) - self._validate_vae_stage_runtime_config(vae_decode_config) + self._validate_vae_stage_runtime_config(vae_encode_config, allow_parallel=False) + self._validate_vae_stage_runtime_config(vae_decode_config, allow_parallel=True) self.vae_encode_device = self._runtime_device(vae_encode_config) self.vae_decode_device = self._runtime_device(vae_decode_config) self.vae_encode_torch_dtype = vae_encode_config.torch_dtype @@ -269,7 +329,6 @@ def init(self, module_manager: ModuleManager, config: LingBotWorldFastPipelineCo "lingbot_world_fast_vae_decode", module_manager, vae_decode_config ) self.vae_encode_worker = ParallelWorker(vae_encode_stage) - self.vae_decode_worker = ParallelWorker(vae_decode_stage) dit_runtime_config = config.dit_config dit_device = "cpu" if dit_runtime_config.parallel_config.world_size > 1 else self.device @@ -284,16 +343,30 @@ def init(self, module_manager: ModuleManager, config: LingBotWorldFastPipelineCo self.dit.set_causal_attention_window(config.local_attn_size, config.sink_size) denoise_stage = LingBotWorldFastDenoisingStage("lingbot_world_fast_denoise", module_manager, dit_runtime_config) + colocate_vae_decode = ( + dit_runtime_config.parallel_config.world_size > 1 + and dit_runtime_config.parallel_config.device_ids == vae_decode_config.parallel_config.device_ids + and dit_runtime_config.parallel_config.world_size == vae_decode_config.parallel_config.world_size + ) + if colocate_vae_decode: + denoise_stage.attach_vae_decode_stage(vae_decode_stage) self.denoise_stage = ( ParallelWorker(denoise_stage) if dit_runtime_config.parallel_config.world_size > 1 else denoise_stage ) + self.vae_decode_worker = ( + _CoLocatedVAEDecodeWorker(self.denoise_stage) if colocate_vae_decode else ParallelWorker(vae_decode_stage) + ) @staticmethod - def _validate_vae_stage_runtime_config(runtime_config: ModelRuntimeConfig) -> None: - """Restrict the VAE stage to one configured GPU until VAE model parallelism exists.""" + def _validate_vae_stage_runtime_config( + runtime_config: ModelRuntimeConfig, + *, + allow_parallel: bool = False, + ) -> None: + """Validate a VAE stage, allowing height-sharded decode workers only.""" runtime_config.parallel_config.validate() - if runtime_config.parallel_config.world_size != 1: - raise ValueError("LingBot VAE stage currently requires exactly one GPU") + if not allow_parallel and runtime_config.parallel_config.world_size != 1: + raise ValueError("LingBot VAE encode stage currently requires exactly one GPU") @staticmethod def _resolve_self_kv_size( @@ -660,6 +733,7 @@ def _create_initialized_session( generator_state=denoise_generator.get_state().tolist(), noise_generator_state=noise_generator.get_state().tolist(), noise_shape=(1, 16, session_config.chunk_size, lat_h, lat_w), + prompt_emb=prompt_emb, timestep_indices=getattr(self.config, "timestep_indices", (0, 179, 358, 679)), ) if isinstance(self.denoise_stage, ParallelWorker): diff --git a/telefuser/pipelines/lingbot_world_fast/service.py b/telefuser/pipelines/lingbot_world_fast/service.py index 376c3bf..d2a9805 100644 --- a/telefuser/pipelines/lingbot_world_fast/service.py +++ b/telefuser/pipelines/lingbot_world_fast/service.py @@ -460,6 +460,9 @@ def discard_first(predicate: Callable[[dict], bool]) -> bool: return True return False + def is_measurement_status(item: dict) -> bool: + return item.get("type") == "status" and isinstance(item.get("measurement"), dict) + if output_queue.full(): discarded = False if payload_type in _VIDEO_OUTPUT_TYPES: @@ -470,7 +473,18 @@ def discard_first(predicate: Callable[[dict], bool]) -> bool: return elif payload_type == "status": stage = payload.get("stage") - discarded = discard_first(lambda item: item.get("type") == "status" and item.get("stage") == stage) + if is_measurement_status(payload): + discarded = discard_first( + lambda item: item.get("type") == "status" and not is_measurement_status(item) + ) + else: + discarded = discard_first( + lambda item: ( + item.get("type") == "status" + and item.get("stage") == stage + and not is_measurement_status(item) + ) + ) if not discarded: with state.metrics_lock: state.dropped_status_payloads += 1 @@ -1187,6 +1201,7 @@ def submit_chunk( applied_controls = controls_by_chunk.pop(result_index, None) control_received_at = control_received_at_by_chunk.pop(result_index, None) chunk_facts = self._finish_benchmark_measurement(measurements_by_chunk.pop(result_index, None)) + chunk_profile = streaming_runtime._pop_chunk_profile(streaming_session, result_index) if state.config.show_control_hud: frames = self._overlay_control_hud(frames, applied_controls) self._put_output( @@ -1245,6 +1260,7 @@ def submit_chunk( "frames": len(frames), "compute_seconds": chunk_facts["seconds"], "memory": chunk_facts["memory"], + "phases": chunk_profile, } } if chunk_facts is not None diff --git a/telefuser/pipelines/lingbot_world_fast/streaming.py b/telefuser/pipelines/lingbot_world_fast/streaming.py index 786c9ce..07cc005 100644 --- a/telefuser/pipelines/lingbot_world_fast/streaming.py +++ b/telefuser/pipelines/lingbot_world_fast/streaming.py @@ -3,8 +3,10 @@ from __future__ import annotations import threading +import time from collections.abc import Callable -from dataclasses import dataclass +from contextlib import nullcontext +from dataclasses import dataclass, field from typing import TYPE_CHECKING import torch @@ -52,6 +54,7 @@ class _LingBotStreamingSessionEntry: progress_callback: Callable[..., None] | None next_condition_index: int = 0 next_control_index: int = 0 + chunk_profiles: dict[int, dict[str, object]] = field(default_factory=dict) class LingBotWorldFastStreamingRuntime: @@ -62,6 +65,8 @@ def __init__(self, pipeline: LingBotWorldFastPipeline) -> None: self._lock = threading.RLock() self._sessions: dict[str, _LingBotStreamingSessionEntry] = {} self._closed = False + self._serialize_dit_decode = self._dit_decode_devices_overlap() + self._dit_decode_lock = threading.Lock() actors = { "encode": ParallelWorkerStageActor( pipeline.vae_encode_worker, @@ -72,12 +77,9 @@ def __init__(self, pipeline: LingBotWorldFastPipeline) -> None: session_closer=self._release_encode_session, ), "denoise": self._denoise_actor(), - "decode": ParallelWorkerStageActor( - pipeline.vae_decode_worker, - "decode_chunk", - self._decode_inputs, - self._decode_outputs, - close_worker=False, + "decode": LocalStageActor( + self._decode, + name="lingbot-decode-actor", session_closer=self._release_decode_session, ), } @@ -251,6 +253,27 @@ def session_metrics(self, session: LingBotWorldFastStreamingSession) -> Streamin self._require_session(session) return self.orchestrator.session_metrics(session.session_id) + def _pop_chunk_profile(self, session: LingBotWorldFastStreamingSession, index: int) -> dict[str, object]: + entry = self._require_session(session) + with self._lock: + profile = entry.chunk_profiles.pop(index, {}) + for stage_id in ("encode", "denoise", "decode"): + timing = next( + ( + item + for item in self.orchestrator.stage_timings(session.session_id, stage_id) + if item.sequence_id == index + ), + None, + ) + if timing is None: + continue + if timing.admitted_at is not None and timing.completed_at is not None: + profile[f"{stage_id}_actor_seconds"] = timing.completed_at - timing.admitted_at + if timing.inputs_ready_at is not None and timing.admitted_at is not None: + profile[f"{stage_id}_queue_seconds"] = timing.admitted_at - timing.inputs_ready_at + return profile + def wait_until_idle(self, session: LingBotWorldFastStreamingSession, timeout: float = 5.0) -> bool: """Wait until the session has no admitted or immediately admissible work.""" self._require_session(session) @@ -327,6 +350,21 @@ def _denoise_actor(self) -> LocalStageActor: session_closer=self._release_denoise_session, ) + @staticmethod + def _runtime_device_ids(runtime_config: object) -> set[int]: + parallel_config = runtime_config.parallel_config + if parallel_config.device_ids is not None: + return set(parallel_config.device_ids) + return {runtime_config.device_id} + + def _dit_decode_devices_overlap(self) -> bool: + config = getattr(self.pipeline, "config", None) + if config is None: + return False + dit_devices = self._runtime_device_ids(config.dit_config) + decode_devices = self._runtime_device_ids(config.vae_decode_config) + return not dit_devices.isdisjoint(decode_devices) + def _entry_for_context(self, context: StreamingSessionContext) -> _LingBotStreamingSessionEntry: with self._lock: try: @@ -428,10 +466,15 @@ def _denoise_kwargs(self, invocation: StreamingStageInvocation) -> dict[str, obj return { "cache_handle": runtime.cache_handle, "condition_chunk": invocation.inputs["condition"], - "prompt_emb": runtime.prompt_emb, + "prompt_emb": None, "control_chunk": invocation.inputs["control"], "current_start": index * runtime.chunk_size * runtime.frame_tokens, "max_attention_size": runtime.max_attention_size, + "_local_vae_handoff": bool( + runtime.world_kv_binding is None + and getattr(self.pipeline.vae_decode_worker, "uses_local_latent_handoff", False) + ), + "_benchmark_profile": runtime.config.benchmark_metrics, } @torch.inference_mode() @@ -449,9 +492,25 @@ def _denoise(self, invocation: StreamingStageInvocation) -> dict[str, object]: else: self.pipeline._notify_progress(entry.progress_callback, "denoising_chunk", index=index) kwargs = self._denoise_kwargs(invocation) - latent = self.pipeline.denoise_stage.denoise_and_update_cache(**kwargs) - if callable(latent): - latent = latent() + lock = self._dit_decode_lock if self._serialize_dit_decode else nullcontext() + lock_started_at = time.perf_counter() + with lock: + worker_started_at = time.perf_counter() + result = self.pipeline.denoise_stage.denoise_and_update_cache(**kwargs) + submit_finished_at = time.perf_counter() + if callable(result): + result = result() + worker_finished_at = time.perf_counter() + if isinstance(result, tuple): + latent, profile = result + profile["denoise_lock_wait_seconds"] = worker_started_at - lock_started_at + profile["denoise_submit_seconds"] = submit_finished_at - worker_started_at + profile["denoise_result_wait_seconds"] = worker_finished_at - submit_finished_at + profile["denoise_worker_seconds"] = worker_finished_at - worker_started_at + with self._lock: + entry.chunk_profiles.setdefault(index, {}).update(profile) + else: + latent = result self.pipeline._notify_progress(entry.progress_callback, "chunk_denoised", index=index) if runtime.world_kv_binding is not None: try: @@ -476,11 +535,16 @@ def _decode_inputs(self, invocation: StreamingStageInvocation) -> tuple[tuple[ob "latents": invocation.inputs["latent"], "is_first_clip": index == 0, "is_last_clip": index == runtime.chunk_count - 1, + "_benchmark_profile": runtime.config.benchmark_metrics, } def _decode_outputs(self, value: torch.Tensor, invocation: StreamingStageInvocation) -> dict[str, object]: entry = self._entry_for_invocation(invocation) - frames = self.pipeline.tensor2video(value) + if value.dtype == torch.uint8: + arrays = value.permute(1, 2, 3, 0).contiguous().numpy() + frames = [Image.fromarray(array) for array in arrays] + else: + frames = self.pipeline.tensor2video(value) self.pipeline._notify_progress( entry.progress_callback, "chunk_decoded", @@ -488,3 +552,31 @@ def _decode_outputs(self, value: torch.Tensor, invocation: StreamingStageInvocat frames=len(frames), ) return {"frames": frames} + + @torch.inference_mode() + def _decode(self, invocation: StreamingStageInvocation) -> dict[str, object]: + args, kwargs = self._decode_inputs(invocation) + lock = self._dit_decode_lock if self._serialize_dit_decode else nullcontext() + lock_started_at = time.perf_counter() + with lock: + worker_started_at = time.perf_counter() + result = self.pipeline.vae_decode_worker.decode_chunk(*args, **kwargs) + submit_finished_at = time.perf_counter() + if callable(result): + result = result() + worker_finished_at = time.perf_counter() + if isinstance(result, tuple): + value, profile = result + profile["decode_lock_wait_seconds"] = worker_started_at - lock_started_at + profile["decode_submit_seconds"] = submit_finished_at - worker_started_at + profile["decode_result_wait_seconds"] = worker_finished_at - submit_finished_at + profile["decode_worker_seconds"] = worker_finished_at - worker_started_at + else: + value, profile = result, {} + convert_start = time.perf_counter() + output = self._decode_outputs(value, invocation) + profile["tensor_to_frames_seconds"] = time.perf_counter() - convert_start + entry = self._entry_for_invocation(invocation) + with self._lock: + entry.chunk_profiles.setdefault(invocation.key.sequence_id, {}).update(profile) + return output diff --git a/telefuser/pipelines/lingbot_world_fast/vae_stage.py b/telefuser/pipelines/lingbot_world_fast/vae_stage.py index b2c3f0b..e15e3d5 100644 --- a/telefuser/pipelines/lingbot_world_fast/vae_stage.py +++ b/telefuser/pipelines/lingbot_world_fast/vae_stage.py @@ -2,6 +2,7 @@ from __future__ import annotations +import time from dataclasses import dataclass, field import torch @@ -14,6 +15,8 @@ WanVideoVAE, WanVideoVAEStreamingDecodeState, WanVideoVAEStreamingEncodeState, + _convert_conv3d_to_channels_last_3d, + _enable_spatial_parallel_decode, ) @@ -102,6 +105,7 @@ class _VAEEncodeCacheState: condition_image: torch.Tensor | None encoder_state: WanVideoVAEStreamingEncodeState = field(default_factory=WanVideoVAEStreamingEncodeState) + latent_condition: torch.Tensor | None = None pool_slot: int | None = None @@ -114,6 +118,9 @@ def __init__(self, name: str, module_manager: ModuleManager, model_runtime_confi if self.vae is None: raise ValueError("LingBot VAE encode stage requires a loaded wan_video_vae module") self.model_names = ["vae"] + # Condition chunks reuse the same bounded shapes for the session lifetime. + # Retain allocator blocks instead of forcing a driver allocation per call. + self.empty_cache_after_call = False self._cache_registry: dict[int, _VAEEncodeCacheState] = {} self._observed_session_cache_bytes = 0 self._cache_layout: dict[int, tuple[torch.dtype, int]] = {} @@ -146,8 +153,6 @@ def configure_cache_pool(self, capacity: int) -> VAECachePoolProfile: raise ValueError(f"cache pool capacity must be positive, got {capacity}") if self._cache_registry: raise RuntimeError("cannot configure the LingBot VAE encode cache pool while sessions are active") - if not self._cache_layout: - raise RuntimeError("cannot configure the LingBot VAE encode cache pool before warmup") existing = self._cache_pool if existing is not None: if existing.capacity != capacity: @@ -194,35 +199,43 @@ def initialize_cache(self, cache_handle: int, condition_image: torch.Tensor) -> def encode_condition_chunk( self, cache_handle: int, chunk_index: int, chunk_count: int, chunk_size: int, height: int, width: int ) -> torch.Tensor: - """Encode one condition chunk and return CPU transport features.""" + """Encode the bounded zero-frame prefix once and return one condition chunk.""" state = self._cache_registry[cache_handle] - is_first = chunk_index == 0 - video = torch.zeros( - (3, 1 + 4 * (chunk_size - 1) if is_first else 4 * chunk_size, height, width), - device=self.device, - dtype=self.torch_dtype, - ) - if is_first: + if state.latent_condition is None: if state.condition_image is None: - raise RuntimeError("The first condition chunk requires the session image tensor") + raise RuntimeError("The first condition request requires the session image tensor") + target_latent_frames = chunk_count * chunk_size + encoded_latent_frames = min(target_latent_frames, 16) + video = torch.zeros( + (3, 1 + 4 * (encoded_latent_frames - 1), height, width), + device=self.device, + dtype=self.torch_dtype, + ) video[:, 0] = state.condition_image - latent = self.vae.cached_encode_withflag( - video, - device=self.device, - is_first_clip=is_first, - is_last_clip=chunk_index == chunk_count - 1, - encode_state=state.encoder_state, - ) - self._observe_cache(state.encoder_state.feat_cache) - if self._cache_pool is not None and state.pool_slot is not None: - self._cache_pool.stabilize(state.encoder_state.feat_cache, state.pool_slot) - if latent.shape[1] != chunk_size: - raise RuntimeError(f"VAE condition chunk has {latent.shape[1]} latent frames, expected {chunk_size}") - mask = torch.zeros((4, chunk_size, latent.shape[2], latent.shape[3]), device=latent.device, dtype=latent.dtype) - if is_first: - mask[:, 0] = 1 + latent = self.vae.cached_encode_withflag( + video, + device=self.device, + is_first_clip=True, + is_last_clip=True, + encode_state=state.encoder_state, + ) + if latent.shape[1] != encoded_latent_frames: + raise RuntimeError( + f"VAE condition prefix has {latent.shape[1]} latent frames, expected {encoded_latent_frames}" + ) + state.latent_condition = latent.cpu() state.condition_image = None - return torch.cat([mask, latent], dim=0).unsqueeze(0).cpu() + + latent_condition = state.latent_condition + start = chunk_index * chunk_size + available = latent_condition[:, start : start + chunk_size] + if available.shape[1] < chunk_size: + tail = latent_condition[:, -1:].expand(-1, chunk_size - available.shape[1], -1, -1) + available = torch.cat([available, tail], dim=1) + mask = torch.zeros((4, chunk_size, available.shape[2], available.shape[3]), dtype=available.dtype) + if chunk_index == 0: + mask[:, 0] = 1 + return torch.cat([mask, available], dim=0).unsqueeze(0) def release_cache(self, cache_handle: int) -> bool: """Release encoder state for one session.""" @@ -244,6 +257,18 @@ class _VAEDecodeCacheState: decoder_state: WanVideoVAEStreamingDecodeState = field(default_factory=WanVideoVAEStreamingDecodeState) pool_slot: int | None = None + output_buffer: torch.Tensor | None = None + exported_output_buffer: torch.Tensor | None = None + + +def _copy_frames_to_shared_cpu(state: _VAEDecodeCacheState, frames: torch.Tensor) -> torch.Tensor: + converted = ((frames + 1) * 127.5).clamp_(0, 255).to(dtype=torch.uint8) + output = state.output_buffer + if output is None or output.shape != converted.shape: + output = torch.empty(converted.shape, dtype=torch.uint8, device="cpu").share_memory_() + state.output_buffer = output + output.copy_(converted) + return output class LingBotWorldFastVAEDecodeStage(BaseStage): @@ -260,6 +285,11 @@ def __init__(self, name: str, module_manager: ModuleManager, model_runtime_confi self._cache_layout: dict[int, tuple[torch.dtype, int]] = {} self._cache_pool: _VAECachePool | None = None + def parallel_models(self) -> None: + """Shard decoder feature maps across the configured VAE worker ranks.""" + _enable_spatial_parallel_decode(self.vae) + _convert_conv3d_to_channels_last_3d(self.vae.model.decoder) + def _observe_cache(self, cache: list[object]) -> None: self._observed_session_cache_bytes = max(self._observed_session_cache_bytes, _cache_tensor_bytes(cache)) for index, item in enumerate(cache): @@ -330,10 +360,20 @@ def initialize_cache(self, cache_handle: int) -> bool: @with_model_offload(["vae"]) def decode_chunk( - self, cache_handle: int, latents: torch.Tensor, is_first_clip: bool, is_last_clip: bool - ) -> torch.Tensor: + self, + cache_handle: int, + latents: torch.Tensor, + is_first_clip: bool, + is_last_clip: bool, + _benchmark_profile: bool = False, + ) -> torch.Tensor | None | tuple[torch.Tensor | None, dict[str, float]]: """Decode one latent chunk and return CPU frame tensors.""" state = self._cache_registry[cache_handle] + decode_start = decode_end = None + if _benchmark_profile and self.device.type == "cuda": + decode_start = torch.cuda.Event(enable_timing=True) + decode_end = torch.cuda.Event(enable_timing=True) + decode_start.record() frames = self.vae.cached_decode_withflag( latents, device=self.device, @@ -341,10 +381,33 @@ def decode_chunk( is_last_clip=is_last_clip, decode_state=state.decoder_state, ) + profile = None + if decode_end is not None: + decode_end.record() + decode_end.synchronize() + profile = {"vae_decode_gpu_seconds": decode_start.elapsed_time(decode_end) / 1000.0} self._observe_cache(state.decoder_state.feat_cache) if self._cache_pool is not None and state.pool_slot is not None: self._cache_pool.stabilize(state.decoder_state.feat_cache, state.pool_slot) - return frames.cpu() + if torch.distributed.is_initialized(): + if torch.distributed.get_rank() != 0: + result = None + else: + transfer_start = time.perf_counter() + shared_result = _copy_frames_to_shared_cpu(state, frames) + if state.exported_output_buffer is shared_result: + result = None + else: + state.exported_output_buffer = shared_result + result = shared_result + if profile is not None: + profile["gpu_to_cpu_seconds"] = time.perf_counter() - transfer_start + else: + transfer_start = time.perf_counter() + result = frames.cpu() + if profile is not None: + profile["gpu_to_cpu_seconds"] = time.perf_counter() - transfer_start + return (result, profile) if profile is not None else result def release_cache(self, cache_handle: int) -> bool: """Release decoder state for one session.""" diff --git a/telefuser/worker/parallel_worker.py b/telefuser/worker/parallel_worker.py index 17b01c3..b19b840 100644 --- a/telefuser/worker/parallel_worker.py +++ b/telefuser/worker/parallel_worker.py @@ -12,6 +12,7 @@ import time from collections.abc import Callable from datetime import timedelta +from multiprocessing.queues import SimpleQueue from queue import Empty from typing import TYPE_CHECKING, Any @@ -52,8 +53,8 @@ def to_device(data: Any, device: str | torch.device) -> Any: def _worker_loop( rank: int, world_size: int, - queue_in: list[mp.Queue], - queue_out: mp.Queue, + queue_in: list[SimpleQueue], + queue_out: SimpleQueue, stage: BaseStage, master_port: int, ) -> None: @@ -67,13 +68,12 @@ def _worker_loop( args = None kwargs = None try: + parallel_config = stage.model_runtime_config.parallel_config + # Avoid host-wide launch pools in every spawned CUDA worker, including + # single-rank workers such as the LingBot condition encoder. + torch.set_num_threads(parallel_config.worker_intra_op_threads) device = stage.device if world_size > 1: - parallel_config = stage.model_runtime_config.parallel_config - # Match torchrun's per-rank default. Letting every spawned worker - # inherit the host-wide intra-op pool oversubscribes CPU launch - # threads and can leave accelerators idle between eager kernels. - torch.set_num_threads(parallel_config.worker_intra_op_threads) os.environ["RANK"] = str(rank) os.environ["WORLD_SIZE"] = str(world_size) os.environ["MASTER_ADDR"] = "localhost" @@ -183,8 +183,11 @@ def __init__( raise RuntimeError("Failed to set start method to spawn:", e) spawn_ctx = mp.get_context("spawn") - self.queue_in: list[mp.Queue] = [spawn_ctx.Queue() for _ in range(self.world_size)] - self.queue_out: mp.Queue = spawn_ctx.Queue() + # Queue uses a background feeder thread for every process. These messages + # only carry small commands and shared-tensor handles, so synchronous + # SimpleQueue writes avoid feeder scheduling tails between GPU stages. + self.queue_in: list[SimpleQueue] = [spawn_ctx.SimpleQueue() for _ in range(self.world_size)] + self.queue_out: SimpleQueue = spawn_ctx.SimpleQueue() master_port = PortAllocator().get_free_port_in_interval() logger.info(f"parallel worker {self.name} with port {master_port}, world_size={self.world_size}") @@ -246,7 +249,12 @@ def _mark_failed(self, reason: str) -> None: def _wait_result(self, method_name: str) -> Any: try: - result = self.queue_out.get(timeout=self.timeout) + # SimpleQueue does not expose a timeout argument. Its reader is the + # underlying multiprocessing Connection and provides the same bounded + # wait without reintroducing a feeder or polling thread. + if not self.queue_out._reader.poll(self.timeout): + raise Empty + result = self.queue_out.get() except Empty as exc: reason = f"{method_name} timeout after {self.timeout} seconds" self._mark_failed(reason) diff --git a/tests/unit/models/test_wan_video_vae_spatial.py b/tests/unit/models/test_wan_video_vae_spatial.py new file mode 100644 index 0000000..b08eafa --- /dev/null +++ b/tests/unit/models/test_wan_video_vae_spatial.py @@ -0,0 +1,62 @@ +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from telefuser.distributed.vae_spatial import _SpatialParallelConv2d +from telefuser.models.wan_video_vae import ( + AttentionBlock, + CausalConv3d, + Decoder3d, + WanVideoVAE, + _SpatialParallelCausalConv3d, + _enable_spatial_parallel_decode, +) + + +def test_enable_spatial_parallel_decode_preserves_parameters_and_is_idempotent() -> None: + decoder = Decoder3d(dim=8, z_dim=4) + vae = SimpleNamespace(model=SimpleNamespace(decoder=decoder)) + parameter_ids = {name: id(parameter) for name, parameter in decoder.named_parameters()} + + with patch("telefuser.models.wan_video_vae._spatial_world_size", return_value=4): + converted = _enable_spatial_parallel_decode(vae) + converted_again = _enable_spatial_parallel_decode(vae) + + assert converted > 0 + assert converted_again == 0 + assert decoder._spatial_parallel is True + assert parameter_ids == {name: id(parameter) for name, parameter in decoder.named_parameters()} + assert all( + isinstance(module, _SpatialParallelCausalConv3d) + for module in decoder.modules() + if isinstance(module, CausalConv3d) + ) + assert any(isinstance(module, _SpatialParallelConv2d) for module in decoder.modules()) + assert all(module._spatial_parallel for module in decoder.modules() if isinstance(module, AttentionBlock)) + + +def test_streaming_spatial_decode_rejects_native_parallelism() -> None: + vae = SimpleNamespace( + model=SimpleNamespace(decoder=Decoder3d(dim=8, z_dim=4)), + parallelism=4, + ) + + with ( + patch("telefuser.models.wan_video_vae._spatial_world_size", return_value=4), + pytest.raises(RuntimeError, match="mutually exclusive"), + ): + _enable_spatial_parallel_decode(vae) + + +def test_native_parallelism_rejects_streaming_spatial_decode() -> None: + vae = SimpleNamespace( + model=SimpleNamespace(decoder=SimpleNamespace(_spatial_parallel=True)), + parallelism=1, + ) + + with pytest.raises(RuntimeError, match="mutually exclusive"): + WanVideoVAE.set_parallelism(vae, 4) + + assert vae.parallelism == 1 + WanVideoVAE.set_parallelism(vae, 1) diff --git a/tests/unit/pipelines/lingbot_world_fast/test_pipeline_lifecycle.py b/tests/unit/pipelines/lingbot_world_fast/test_pipeline_lifecycle.py index f8971e2..e5d3802 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_pipeline_lifecycle.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_pipeline_lifecycle.py @@ -1,8 +1,9 @@ from unittest.mock import MagicMock +import torch from PIL import Image -from telefuser.pipelines.lingbot_world_fast.pipeline import LingBotWorldFastPipeline +from telefuser.pipelines.lingbot_world_fast.pipeline import LingBotWorldFastPipeline, _CoLocatedVAEDecodeWorker from telefuser.pipelines.lingbot_world_fast.session import ( LingBotWorldFastGenerationSession, LingBotWorldFastSessionConfig, @@ -42,3 +43,64 @@ def test_pipeline_close_delegates_to_parallel_worker() -> None: pipeline.close() worker.close.assert_called_once_with() + + +def test_colocated_vae_decode_worker_maps_lifecycle_calls() -> None: + worker = MagicMock() + worker.name = "denoise" + proxy = _CoLocatedVAEDecodeWorker(worker) + + assert proxy.uses_local_latent_handoff is True + + proxy.reset_device_memory_peak() + proxy.device_memory_snapshots() + proxy.estimate_session_cache_bytes() + proxy.observed_session_cache_bytes() + proxy.configure_cache_pool(3) + proxy.initialize_cache(7, sync=True) + proxy.decode_chunk(7, "latent") + proxy.release_cache(7, sync=True) + + worker.reset_vae_decode_device_memory_peak.assert_called_once_with(sync=True) + worker.vae_decode_device_memory_snapshots.assert_called_once_with(sync=True) + worker.estimate_vae_decode_session_cache_bytes.assert_called_once_with(sync=True) + worker.observed_vae_decode_session_cache_bytes.assert_called_once_with(sync=True) + worker.configure_vae_decode_cache_pool.assert_called_once_with(3, sync=True) + worker.initialize_vae_decode_cache.assert_called_once_with(7, sync=True) + worker.decode_chunk.assert_called_once_with(7, "latent") + worker.release_vae_decode_cache.assert_called_once_with(7, sync=True) + + +def test_colocated_vae_decode_worker_reuses_registered_output_buffer() -> None: + worker = MagicMock() + worker.name = "denoise" + first = torch.ones(3, 4, 2, 2, dtype=torch.uint8).share_memory_() + worker.decode_chunk.side_effect = [ + lambda: (first, {"chunk": 0}), + lambda: (None, {"chunk": 1}), + ] + worker.release_vae_decode_cache.return_value = True + proxy = _CoLocatedVAEDecodeWorker(worker) + + first_output, first_profile = proxy.decode_chunk(7, "latent")() + second_output, second_profile = proxy.decode_chunk(7, "latent")() + proxy.release_cache(7, sync=True) + + assert first_output.untyped_storage().data_ptr() == first.untyped_storage().data_ptr() + assert second_output.untyped_storage().data_ptr() == first.untyped_storage().data_ptr() + assert first_profile == {"chunk": 0} + assert second_profile == {"chunk": 1} + assert proxy._output_buffers == {} + + +def test_colocated_vae_decode_worker_replaces_resized_output_buffer() -> None: + worker = MagicMock() + worker.name = "denoise" + first = torch.ones(3, 13, 2, 2, dtype=torch.uint8).share_memory_() + resized = torch.ones(3, 16, 2, 2, dtype=torch.uint8).share_memory_() + worker.decode_chunk.side_effect = [lambda: first, lambda: resized, lambda: None] + proxy = _CoLocatedVAEDecodeWorker(worker) + + assert proxy.decode_chunk(7, "latent")().shape[1] == 13 + assert proxy.decode_chunk(7, "latent")().shape[1] == 16 + assert proxy.decode_chunk(7, "latent")().shape[1] == 16 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 dea980d..89a80ca 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 @@ -779,6 +779,21 @@ def test_output_queue_discards_stale_video_and_records_runtime_metrics() -> None assert LingBotWorldFastService._runtime_metrics(state)["output_queue_high_watermark"] == 2 +def test_output_queue_preserves_benchmark_measurements_under_backpressure() -> None: + state = _state() + state.output_queue = asyncio.Queue(maxsize=2) + first = {"type": "status", "stage": "chunk_sent", "measurement": {"index": 0}} + second = {"type": "status", "stage": "chunk_sent", "measurement": {"index": 1}} + + LingBotWorldFastService._enqueue_output(state, first) + LingBotWorldFastService._enqueue_output(state, {"type": "status", "stage": "generating_chunk"}) + LingBotWorldFastService._enqueue_output(state, second) + LingBotWorldFastService._enqueue_output(state, {"type": "status", "stage": "chunk_sent"}) + + assert list(state.output_queue._queue) == [first, second] + assert LingBotWorldFastService._runtime_metrics(state)["dropped_status_payloads"] == 2 + + def test_stream_progress_reports_duration_frames_and_chunks() -> None: service = LingBotWorldFastService(MagicMock(), max_generation_seconds=20.0) state = _state() diff --git a/tests/unit/pipelines/lingbot_world_fast/test_stream_example.py b/tests/unit/pipelines/lingbot_world_fast/test_stream_example.py index 241da03..9b67d66 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_stream_example.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_stream_example.py @@ -63,8 +63,9 @@ def test_unified_example_get_pipeline_maps_ppl_config_to_internal_workers() -> N assert config.dit_config.parallel_config.enable_fsdp is offline_example.PPL_CONFIG["enable_fsdp"] assert config.vae_encode_config.device_id == 0 assert config.vae_encode_config.parallel_config.device_ids == [0] - assert config.vae_decode_config.device_id == 1 - assert config.vae_decode_config.parallel_config.device_ids == [1] + assert config.vae_decode_config.device_id == 0 + assert config.vae_decode_config.parallel_config.device_ids == [0, 1, 2, 3] + assert config.vae_decode_config.parallel_config.sp_ulysses_degree == 4 load_calls = module_manager.load_model.call_args_list assert [call.args[0] for call in load_calls[:2]] == [ "/models/Wan2.2-I2V-A14B/Wan2.1_VAE.pth", @@ -76,10 +77,32 @@ def test_unified_example_get_pipeline_maps_ppl_config_to_internal_workers() -> N def test_unified_example_resolves_fixed_gpu_layouts() -> None: - assert offline_example._resolve_stage_devices(2) == ([0, 1], 0, 1) - assert offline_example._resolve_stage_devices(4) == ([0, 1, 2, 3], 0, 1) - assert offline_example._resolve_stage_devices(5) == ([0, 1, 2, 3], 4, 4) - assert offline_example._resolve_stage_devices(6) == ([0, 1, 2, 3, 4], 5, 5) + assert offline_example._resolve_stage_devices(2) == ([0, 1], 0, [1]) + assert offline_example._resolve_stage_devices(4) == ([0, 1, 2, 3], 0, [0, 1, 2, 3]) + assert offline_example._resolve_stage_devices(5) == ([0, 1, 2, 3], 4, [4]) + assert offline_example._resolve_stage_devices(6) == ([0, 1, 2, 3, 4], 5, [5]) + + +def test_unified_example_six_gpu_uses_five_dit_gpus_and_one_vae_gpu() -> None: + pipeline = MagicMock() + module_manager = MagicMock() + + with ( + patch.object(offline_example, "ModuleManager", return_value=module_manager), + patch.object(offline_example, "LingBotWorldFastPipeline", return_value=pipeline), + ): + offline_example.get_pipeline( + parallelism=6, + model_root="/models/Wan2.2-I2V-A14B", + fast_model_root="/models/lingbot-world-fast", + ) + + config = pipeline.init.call_args.args[1] + assert config.dit_config.parallel_config.device_ids == [0, 1, 2, 3, 4] + assert config.dit_config.parallel_config.sp_ulysses_degree == 5 + assert config.vae_encode_config.parallel_config.device_ids == [5] + assert config.vae_decode_config.parallel_config.device_ids == [5] + assert config.vae_decode_config.parallel_config.sp_ulysses_degree == 1 def test_unified_example_get_service_uses_passed_gpu_num_and_ppl_fps() -> None: diff --git a/tests/unit/pipelines/lingbot_world_fast/test_streaming.py b/tests/unit/pipelines/lingbot_world_fast/test_streaming.py index e309ac8..955d6ca 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_streaming.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_streaming.py @@ -127,6 +127,19 @@ def test_independent_vae_stage_configs_are_resolved_independently() -> None: LingBotWorldFastPipeline._validate_vae_stage_runtime_config(config.vae_decode_config) +def test_only_vae_decode_stage_accepts_spatial_parallel_workers() -> None: + config = ModelRuntimeConfig( + device_type="cuda", + device_id=0, + torch_dtype=torch.float32, + parallel_config=ParallelConfig(device_ids=[0, 1, 2, 3], sp_ulysses_degree=4), + ) + + LingBotWorldFastPipeline._validate_vae_stage_runtime_config(config, allow_parallel=True) + with pytest.raises(ValueError, match="VAE encode stage currently requires exactly one GPU"): + LingBotWorldFastPipeline._validate_vae_stage_runtime_config(config) + + def test_streaming_session_routes_one_chunk_through_three_stages() -> None: pipeline = _Pipeline() runtime = LingBotWorldFastGenerationSession( diff --git a/tests/unit/pipelines/lingbot_world_fast/test_vae_stage_capacity.py b/tests/unit/pipelines/lingbot_world_fast/test_vae_stage_capacity.py index ee92243..9988fa3 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_vae_stage_capacity.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_vae_stage_capacity.py @@ -1,10 +1,32 @@ from __future__ import annotations +from types import SimpleNamespace +from unittest.mock import patch + import pytest import torch from telefuser.pipelines.lingbot_world_fast import vae_stage -from telefuser.pipelines.lingbot_world_fast.vae_stage import _VAECachePool, _cache_tensor_bytes +from telefuser.pipelines.lingbot_world_fast.vae_stage import ( + _VAECachePool, + _VAEDecodeCacheState, + _cache_tensor_bytes, + _copy_frames_to_shared_cpu, +) + + +class _RecordingEncoder: + def __init__(self) -> None: + self.frame_counts: list[int] = [] + + def cached_encode_withflag(self, video, device, is_first_clip, is_last_clip, encode_state): + del device, encode_state + assert is_first_clip is True + assert is_last_clip is True + self.frame_counts.append(video.shape[1]) + latent_frames = (video.shape[1] - 1) // 4 + 1 + values = torch.arange(latent_frames, dtype=video.dtype).view(1, latent_frames, 1, 1) + return values.expand(16, -1, 2, 2).clone() def test_cache_tensor_bytes_ignores_non_tensor_entries() -> None: @@ -18,6 +40,16 @@ def test_cache_tensor_bytes_ignores_non_tensor_entries() -> None: assert _cache_tensor_bytes(cache) == 32 +def test_vae_decode_reuses_shared_cpu_output_buffer() -> None: + state = _VAEDecodeCacheState() + first = _copy_frames_to_shared_cpu(state, torch.zeros(3, 4, 2, 2)) + second = _copy_frames_to_shared_cpu(state, torch.ones(3, 4, 2, 2)) + + assert first.is_shared() + assert second.untyped_storage().data_ptr() == first.untyped_storage().data_ptr() + assert torch.equal(second, torch.full_like(second, 255)) + + def test_vae_cache_pool_stabilizes_and_reuses_fixed_slots() -> None: pool = _VAECachePool( capacity=2, @@ -75,3 +107,39 @@ def test_vae_decode_stage_reports_capacity_without_raising() -> None: assert stage.release_cache(1) is True assert stage.initialize_cache(2) is True assert stage.release_cache(2) is True + + +def test_vae_decode_parallelization_converts_decoder_only() -> None: + stage = vae_stage.LingBotWorldFastVAEDecodeStage.__new__(vae_stage.LingBotWorldFastVAEDecodeStage) + decoder = object() + stage.vae = SimpleNamespace(model=SimpleNamespace(decoder=decoder)) + + with ( + patch.object(vae_stage, "_enable_spatial_parallel_decode") as enable_spatial, + patch.object(vae_stage, "_convert_conv3d_to_channels_last_3d") as convert_channels_last, + ): + stage.parallel_models() + + enable_spatial.assert_called_once_with(stage.vae) + convert_channels_last.assert_called_once_with(decoder) + + +def test_vae_encode_stage_encodes_bounded_prefix_once_and_repeats_tail() -> None: + stage = vae_stage.LingBotWorldFastVAEEncodeStage.__new__(vae_stage.LingBotWorldFastVAEEncodeStage) + stage.device = torch.device("cpu") + stage.torch_dtype = torch.float32 + stage.vae = _RecordingEncoder() + stage._cache_registry = {} + stage._cache_pool = None + assert stage.initialize_cache(1, torch.ones(3, 2, 2)) is True + + encode = vae_stage.LingBotWorldFastVAEEncodeStage.encode_condition_chunk.__wrapped__ + first = encode(stage, 1, 0, 5, 4, 2, 2) + tail = encode(stage, 1, 4, 5, 4, 2, 2) + + assert stage.vae.frame_counts == [61] + assert first.shape == (1, 20, 4, 2, 2) + assert torch.equal(first[0, :4, 0], torch.ones(4, 2, 2)) + assert torch.count_nonzero(first[0, :4, 1:]) == 0 + assert torch.count_nonzero(tail[0, :4]) == 0 + assert torch.equal(tail[0, 4:], torch.full((16, 4, 2, 2), 15.0)) diff --git a/tests/unit/pipelines/lingbot_world_v2/test_service.py b/tests/unit/pipelines/lingbot_world_v2/test_service.py index 80e78d0..051c75a 100644 --- a/tests/unit/pipelines/lingbot_world_v2/test_service.py +++ b/tests/unit/pipelines/lingbot_world_v2/test_service.py @@ -28,16 +28,18 @@ def test_v2_unified_example_get_pipeline_maps_ppl_config_to_internal_workers() - pipeline_cls.assert_called_once_with(device="cuda", torch_dtype=torch.bfloat16) assert pipeline.init.call_args.args[0] is module_manager config = pipeline.init.call_args.args[1] + assert config.max_area == 832 * 480 + 1 assert config.local_attn_size == 18 assert config.sink_size == 6 assert config.timestep_indices == (0, 250, 500, 750) - assert config.dit_config.attention_config.attn_impl == AttnImplType.SAGE_ATTN_2_8_8_SM90 - assert config.dit_config.compile_config.enabled is True + assert config.dit_config.attention_config.attn_impl == offline_example.PPL_CONFIG["attn_impl"] + assert config.dit_config.compile_config.enabled is False assert config.dit_config.parallel_config.device_ids == [0, 1, 2, 3] assert config.vae_encode_config.device_id == 0 assert config.vae_encode_config.parallel_config.device_ids == [0] - assert config.vae_decode_config.device_id == 1 - assert config.vae_decode_config.parallel_config.device_ids == [1] + assert config.vae_decode_config.device_id == 0 + assert config.vae_decode_config.parallel_config.device_ids == [0, 1, 2, 3] + assert config.vae_decode_config.parallel_config.sp_ulysses_degree == 4 load_calls = module_manager.load_model.call_args_list assert [call.args[0] for call in load_calls[:2]] == [ "/models/Wan2.2-I2V-A14B/Wan2.1_VAE.pth", @@ -49,7 +51,7 @@ def test_v2_unified_example_get_pipeline_maps_ppl_config_to_internal_workers() - ] -def test_v2_offline_multi_gpu_defaults_to_the_shared_vae_worker() -> None: +def test_v2_offline_four_gpu_uses_spatial_parallel_vae_decode() -> None: pipeline = MagicMock() module_manager = MagicMock() @@ -68,15 +70,38 @@ def test_v2_offline_multi_gpu_defaults_to_the_shared_vae_worker() -> None: assert config.dit_config.parallel_config.sp_ulysses_degree == 4 assert config.vae_encode_config.device_id == 0 assert config.vae_encode_config.parallel_config.device_ids == [0] - assert config.vae_decode_config.device_id == 1 - assert config.vae_decode_config.parallel_config.device_ids == [1] + assert config.vae_decode_config.device_id == 0 + assert config.vae_decode_config.parallel_config.device_ids == [0, 1, 2, 3] + assert config.vae_decode_config.parallel_config.sp_ulysses_degree == 4 def test_v2_unified_example_resolves_fixed_gpu_layouts() -> None: - assert offline_example._resolve_stage_devices(2) == ([0, 1], 0, 1) - assert offline_example._resolve_stage_devices(4) == ([0, 1, 2, 3], 0, 1) - assert offline_example._resolve_stage_devices(5) == ([0, 1, 2, 3], 4, 4) - assert offline_example._resolve_stage_devices(6) == ([0, 1, 2, 3, 4], 5, 5) + assert offline_example._resolve_stage_devices(2) == ([0, 1], 0, [1]) + assert offline_example._resolve_stage_devices(4) == ([0, 1, 2, 3], 0, [0, 1, 2, 3]) + assert offline_example._resolve_stage_devices(5) == ([0, 1, 2, 3], 4, [4]) + assert offline_example._resolve_stage_devices(6) == ([0, 1, 2, 3, 4], 5, [5]) + + +def test_v2_offline_six_gpu_uses_five_dit_gpus_and_one_vae_gpu() -> None: + pipeline = MagicMock() + module_manager = MagicMock() + + with ( + patch.object(offline_example, "ModuleManager", return_value=module_manager), + patch.object(offline_example, "LingBotWorldV2Pipeline", return_value=pipeline), + ): + offline_example.get_pipeline( + parallelism=6, + model_root="/models/Wan2.2-I2V-A14B", + v2_model_root="/models/lingbot-world-v2-14b-causal-fast/transformers", + ) + + config = pipeline.init.call_args.args[1] + assert config.dit_config.parallel_config.device_ids == [0, 1, 2, 3, 4] + assert config.dit_config.parallel_config.sp_ulysses_degree == 5 + assert config.vae_encode_config.parallel_config.device_ids == [5] + assert config.vae_decode_config.parallel_config.device_ids == [5] + assert config.vae_decode_config.parallel_config.sp_ulysses_degree == 1 def test_v2_unified_example_uses_ppl_configured_vae_stage_devices_independently() -> None: @@ -114,6 +139,7 @@ def test_v2_unified_example_service_constructs_v2_session_from_ppl_config() -> N assert isinstance(session_config, LingBotWorldFastSessionConfig) assert session_config.frame_num == 1917 + assert session_config.image.size == (832, 480) assert session_config.chunk_size == 4 assert session_config.frame_policy == "truncate" assert session_config.sample_shift == 10.0 diff --git a/tests/unit/worker/test_parallel_worker.py b/tests/unit/worker/test_parallel_worker.py index ca9800d..f3a0596 100644 --- a/tests/unit/worker/test_parallel_worker.py +++ b/tests/unit/worker/test_parallel_worker.py @@ -127,7 +127,7 @@ def test_put_data_with_queue_cpu(self, mock_port_allocator, mock_mp): mock_spawn_ctx = MagicMock() mock_mp.get_context.return_value = mock_spawn_ctx mock_queue = MagicMock() - mock_spawn_ctx.Queue.return_value = mock_queue + mock_spawn_ctx.SimpleQueue.return_value = mock_queue mock_stage = MagicMock() mock_stage.name = "TestStage" @@ -154,7 +154,8 @@ def test_call_with_sync(self, mock_port_allocator, mock_mp): mock_mp.get_context.return_value = mock_spawn_ctx mock_queue_out = MagicMock() mock_queue_out.get.return_value = torch.randn(2, 3) - mock_spawn_ctx.Queue.return_value = mock_queue_out + mock_spawn_ctx.SimpleQueue.return_value = mock_queue_out + mock_queue_out._reader.poll.return_value = True mock_stage = MagicMock() mock_stage.name = "TestStage" @@ -181,7 +182,8 @@ def test_call_without_sync(self, mock_port_allocator, mock_mp): mock_mp.get_context.return_value = mock_spawn_ctx mock_queue_out = MagicMock() mock_queue_out.get.return_value = torch.randn(2, 3) - mock_spawn_ctx.Queue.return_value = mock_queue_out + mock_spawn_ctx.SimpleQueue.return_value = mock_queue_out + mock_queue_out._reader.poll.return_value = True mock_stage = MagicMock() mock_stage.name = "TestStage" @@ -209,7 +211,8 @@ def test_getattr_arbitrary_method(self, mock_port_allocator, mock_mp): mock_mp.get_context.return_value = mock_spawn_ctx mock_queue_out = MagicMock() mock_queue_out.get.return_value = "result" - mock_spawn_ctx.Queue.return_value = mock_queue_out + mock_spawn_ctx.SimpleQueue.return_value = mock_queue_out + mock_queue_out._reader.poll.return_value = True mock_stage = MagicMock() mock_stage.name = "TestStage" @@ -235,7 +238,8 @@ def test_queue_out_exception(self, mock_port_allocator, mock_mp): mock_mp.get_context.return_value = mock_spawn_ctx mock_queue_out = MagicMock() mock_queue_out.get.return_value = RuntimeError("Worker error") - mock_spawn_ctx.Queue.return_value = mock_queue_out + mock_spawn_ctx.SimpleQueue.return_value = mock_queue_out + mock_queue_out._reader.poll.return_value = True mock_stage = MagicMock() mock_stage.name = "TestStage" @@ -260,8 +264,8 @@ def test_queue_timeout(self, mock_port_allocator, mock_mp): mock_spawn_ctx = MagicMock() mock_mp.get_context.return_value = mock_spawn_ctx mock_queue_out = MagicMock() - mock_queue_out.get.side_effect = Empty() - mock_spawn_ctx.Queue.return_value = mock_queue_out + mock_queue_out._reader.poll.return_value = False + mock_spawn_ctx.SimpleQueue.return_value = mock_queue_out mock_stage = MagicMock() mock_stage.name = "TestStage" @@ -302,6 +306,7 @@ def test_worker_loop_single_process(self, mock_platform, mock_dist): mock_stage = MagicMock() mock_stage.device = "cpu" + mock_stage.model_runtime_config.parallel_config.worker_intra_op_threads = 1 mock_stage.test_method.return_value = torch.randn(2, 3) _worker_loop( @@ -332,6 +337,7 @@ def test_worker_loop_missing_method(self, mock_platform, mock_dist): mock_stage = MagicMock() mock_stage.device = "cpu" + mock_stage.model_runtime_config.parallel_config.worker_intra_op_threads = 1 mock_stage.name = "TestStage" del mock_stage.nonexistent_method # Ensure method doesn't exist @@ -363,6 +369,7 @@ def test_worker_loop_exception_handling(self, mock_platform, mock_dist): mock_stage = MagicMock() mock_stage.device = "cpu" + mock_stage.model_runtime_config.parallel_config.worker_intra_op_threads = 1 mock_stage.test_method.side_effect = RuntimeError("Method error") _worker_loop( From 663c385b179012c5c3de613212d10e8e6eac5f5d Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Sun, 2 Aug 2026 06:27:33 +0000 Subject: [PATCH 04/11] test(aiperf): standardize LingBot stream profiling Record decoded-ready, LiveKit publish, and client metadata timing separately, wait for the controller before pipeline startup, and finish adapter sessions from target completion status. Keep the reproducible TeleFuser and SGLang benchmark launchers, add a direct-output validation path, and document the H100 environment, workload semantics, and compute-versus-delivery metric boundary without embedding comparison results. --- benchmarks/telefuser_aiperf/README.md | 11 +- .../run_sglang_lingbot_world_v2_4gpu.sh | 7 +- .../telefuser_aiperf/adapter.py | 39 +++- .../tests/test_livekit_adapter.py | 21 +++ docs/en/benchmark_aiperf.md | 39 +--- docs/zh/benchmark_aiperf.md | 46 ++--- examples/lingbot/README.md | 118 +++++++----- telefuser/service/livekit/room_client.py | 27 ++- telefuser/service/livekit/worker.py | 16 ++ .../livekit/test_multi_session_worker.py | 3 + .../unit/service/livekit/test_room_client.py | 14 +- tests/unit/service/livekit/test_worker.py | 59 ++++++ .../benchmark_lingbot_world_v2_direct.py | 174 ++++++++++++++++++ 13 files changed, 451 insertions(+), 123 deletions(-) create mode 100644 tools/validation/benchmark_lingbot_world_v2_direct.py diff --git a/benchmarks/telefuser_aiperf/README.md b/benchmarks/telefuser_aiperf/README.md index 9b3e222..ae26164 100644 --- a/benchmarks/telefuser_aiperf/README.md +++ b/benchmarks/telefuser_aiperf/README.md @@ -143,9 +143,9 @@ In terminal 3, run the one-minute LingBot-World v2 workload: bash benchmarks/telefuser_aiperf/scripts/run_stream_bench.sh ``` -The v2 workload is the launcher default and the only checked-in stream workload. It requests 59.75 seconds of media -using the model's fixed attention window. Its AIPerf active window is 240 seconds, so a successful command normally -takes about four minutes rather than one minute. Do not terminate it after media generation becomes quiet. +The v2 workload is the default TeleFuser stream workload. It requests 59.75 seconds of media using the model's fixed +attention window. The 240-second AIPerf active window is a timeout ceiling; a successful run exits after the target +emits its completion status and normally takes about 66 seconds after admission, excluding model loading. A successful run prints `Stream profile sessions: 1/1 succeeded`, an artifact directory, and an HTML report path. Results are written below: @@ -173,9 +173,8 @@ bash benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh ``` The launcher defaults to GPUs `0,1,2,3`, port `30000`, and model -`robbyant/lingbot-world-v2-14b-causal-fast-diffusers`. It uses an installed `sglang` command when available. -For the checked-in `work_dirs/sglang` source, set `SGLANG_PYTHON` to an environment installed with that source's -dependencies. Override the defaults when needed: +`robbyant/lingbot-world-v2-14b-causal-fast-diffusers`. Explicit `SGLANG_SOURCE_DIR` and `SGLANG_PYTHON` values take +precedence over an installed `sglang` command. Override the defaults when needed: ```bash SGLANG_BIN=/path/to/sglang \ diff --git a/benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh b/benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh index 521938d..60fcc16 100755 --- a/benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh +++ b/benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh @@ -19,14 +19,14 @@ if [[ -n "${SGLANG_BIN}" ]]; then exit 1 fi sglang_command=("${SGLANG_BIN}") -elif command -v sglang >/dev/null 2>&1; then - sglang_command=("$(command -v sglang)") elif [[ -n "${SGLANG_PYTHON}" && -x "${SGLANG_PYTHON}" \ && -f "${SGLANG_SOURCE_DIR}/python/sglang/cli/main.py" ]]; then export PYTHONPATH="${SGLANG_SOURCE_DIR}/python${PYTHONPATH:+:${PYTHONPATH}}" sglang_command=("${SGLANG_PYTHON}" "-c" "from sglang.cli.main import main; main()") +elif command -v sglang >/dev/null 2>&1; then + sglang_command=("$(command -v sglang)") else - echo "SGLang is unavailable. Set SGLANG_BIN, or set SGLANG_SOURCE_DIR and a compatible SGLANG_PYTHON." >&2 + echo "SGLang is unavailable. Set SGLANG_SOURCE_DIR and SGLANG_PYTHON, or set SGLANG_BIN." >&2 exit 1 fi @@ -40,6 +40,7 @@ export CUDA_VISIBLE_DEVICES="${SGLANG_CUDA_VISIBLE_DEVICES}" export SGLANG_LINGBOT_LAZY_VAE_ENCODE_BLACK_FRAMES="${SGLANG_LINGBOT_LAZY_VAE_ENCODE_BLACK_FRAMES:-60}" exec "${sglang_command[@]}" serve \ + --model-type diffusion \ --model-path "${SGLANG_MODEL_PATH}" \ --pipeline-class-name LingBotWorldCausalDMDPipeline \ --host "${SGLANG_HOST}" \ diff --git a/benchmarks/telefuser_aiperf/telefuser_aiperf/adapter.py b/benchmarks/telefuser_aiperf/telefuser_aiperf/adapter.py index debf387..63bdab9 100644 --- a/benchmarks/telefuser_aiperf/telefuser_aiperf/adapter.py +++ b/benchmarks/telefuser_aiperf/telefuser_aiperf/adapter.py @@ -150,6 +150,7 @@ def _handle_data_message( self.events.record("done_message", topic=topic) return data = payload.get("data") if isinstance(payload.get("data"), dict) else payload + self._record_transport_profile(data) if data.get("type") == "error" or payload.get("error"): error = data.get("error") or payload.get("error") self.result.error = redact_string(str(error)) @@ -177,11 +178,25 @@ def _handle_status_stage( self.target_ready = True self._try_start_active_window() if stage in {"runtime_ready", "chunk_sent"}: + measurement = data.get("measurement") + normalized_data = data + if stage == "chunk_sent" and isinstance(measurement, Mapping): + phases = measurement.get("phases") + if isinstance(phases, Mapping): + self.events.record( + "target_phase_profile", + chunk_index=measurement.get("index"), + phases=dict(phases), + ) + normalized_data = { + **data, + "measurement": {key: value for key, value in measurement.items() if key != "phases"}, + } record_target_measurement( result=self.result, events=self.events, stage=stage, - data=data, + data=normalized_data, ) if not self.pending_control_acks: return @@ -195,6 +210,28 @@ def _handle_status_stage( 0.0, ) + def _record_transport_profile(self, data: Mapping[str, Any]) -> None: + transport = data.get("transport_measurement") + if not isinstance(transport, Mapping): + return + received_at = time.time() + chunk_index = data.get("index") + measurement = data.get("measurement") + if chunk_index is None and isinstance(measurement, Mapping): + chunk_index = measurement.get("index") + publish_finished_at = transport.get("publish_finished_at") + self.events.record( + "target_transport_profile", + chunk_index=chunk_index, + transport=dict(transport), + client_metadata_received_at=received_at, + publish_to_client_metadata_seconds=( + max(received_at - float(publish_finished_at), 0.0) + if isinstance(publish_finished_at, int | float) + else None + ), + ) + async def _send_control_trace(self) -> None: active_started_at = self._active_start() for event_index, entry in enumerate(self.plan.control_trace): diff --git a/benchmarks/telefuser_aiperf/tests/test_livekit_adapter.py b/benchmarks/telefuser_aiperf/tests/test_livekit_adapter.py index 5766e93..bd18913 100644 --- a/benchmarks/telefuser_aiperf/tests/test_livekit_adapter.py +++ b/benchmarks/telefuser_aiperf/tests/test_livekit_adapter.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +import time from collections.abc import Callable, Mapping from pathlib import Path from typing import Any @@ -108,6 +109,18 @@ async def connect( "compute_seconds": 0.5, "encode_seconds": 0.1, "memory": [], + "phases": { + "dmd_step_seconds": [0.1, 0.1, 0.1, 0.1], + "vae_decode_gpu_seconds": 0.08, + }, + }, + "transport_measurement": { + "decoded_ready_at": time.time() - 1.0, + "publish_started_at": time.time() - 0.5, + "publish_finished_at": time.time() - 0.01, + "publish_seconds": 0.49, + "frames": 3, + "pacing": "realtime", }, }, } @@ -317,6 +330,14 @@ async def test_adapter_normalizes_room_events_and_deletes_target( "http://127.0.0.1:30000/sessions/livekit-session", ) in http.requests assert Path(result.artifacts_event_file or "").is_file() + events = [orjson.loads(line) for line in Path(result.artifacts_event_file or "").read_bytes().splitlines()] + phase_profile = next(item for item in events if item["event"] == "target_phase_profile") + assert phase_profile["chunk_index"] == 0 + assert phase_profile["phases"]["vae_decode_gpu_seconds"] == 0.08 + transport_profile = next(item for item in events if item["event"] == "target_transport_profile") + assert transport_profile["chunk_index"] == 0 + assert transport_profile["transport"]["pacing"] == "realtime" + assert transport_profile["publish_to_client_metadata_seconds"] >= 0.0 @pytest.mark.asyncio diff --git a/docs/en/benchmark_aiperf.md b/docs/en/benchmark_aiperf.md index a703129..f474395 100644 --- a/docs/en/benchmark_aiperf.md +++ b/docs/en/benchmark_aiperf.md @@ -56,10 +56,10 @@ An idle service reports `"livekit_connected":false`; that is expected. Run the b bash benchmarks/telefuser_aiperf/scripts/run_stream_bench.sh ``` -The request contains 59.75 seconds of media, but the configured AIPerf active window is 240 seconds. Allow about four -minutes for the command to finish. Success is `Stream profile sessions: 1/1 succeeded`; reports are created under -`artifacts/telefuser_aiperf/stream_lingbot_v2_1min/`. See the canonical README for model file layout, manual Python -environment selection, history setup, and troubleshooting. +The request contains 59.75 seconds of media. The 240-second active window is a timeout ceiling; a successful run +exits after the target completion status and normally takes about 66 seconds after admission. Success is +`Stream profile sessions: 1/1 succeeded`; reports are created under +`artifacts/telefuser_aiperf/stream_lingbot_v2_1min/`. To profile SGLang on the same four-GPU workload instead, start its server and select the SGLang config: @@ -102,37 +102,6 @@ Target facts follow these rules: Client delivery, target pipeline residence, target phase time, and resource utilization remain separate dimensions. Fields without equivalent semantics remain private or unavailable instead of being forced into a common metric. -## Validated one-minute LingBot-World v2 replay - -The `stream_lingbot_world_v2_1min.json` workload was validated on 2026-07-28 with four H100 80 GB GPUs, BF16 DiT, -FP32 VAE, SageAttention SM90, `torch.compile` enabled, FSDP disabled, `chunk_size=4`, and 16 FPS. A 60-second request -is truncated to 60 complete latent chunks: 957 output frames representing 59.75 seconds of media. LingBot-World v2 -used its fixed `local_attn_size=18` and `sink_size=6` window, so the 240 latent-frame request retained a fixed -27,144-token KV capacity rather than a duration-sized global KV cache. - -| Measurement | Result | -|---|---:| -| Successful sessions | 1 / 1 | -| Target chunks / generated frames | 60 / 957 | -| Client frames received | 947 | -| Target steady chunks / frames | 59 / 944 | -| Configured session runtime | 242.255 s | -| First-frame latency | 6,252.773 ms | -| Client delivery rate (`stream_fps`) | 9.397 FPS | -| Weighted steady chunk compute rate | 4.708 FPS | -| Chunk pipeline residence, mean / p99 | 3.399 / 3.901 s | - -The target emitted all 957 frames and the LiveKit client received 947. AIPerf excluded the first target chunk from -steady-state compute aggregation, leaving 944 generated frames across 59 chunks. The 242-second session runtime is -the configured 240-second active-window limit plus connection overhead; it is not the generation time of the -59.75-second media payload. - -Chunk pipeline residence spans actor admission through output and includes time shared by overlapping encode, DiT, -and decode work. It is therefore expected to exceed adjacent delivery intervals; the report's 4.708 weighted -steady chunk compute FPS must not be read as client stream throughput. The Git-installed AIPerf run exited with code -0. The replay artifact is -`artifacts/telefuser_aiperf/stream_lingbot_v2_1min/20260728_083948_fc4344ba/stream_report.html`. - ## Reproducibility Every result should retain the TeleFuser commit and AIPerf package version, model revision, accelerator model/count, diff --git a/docs/zh/benchmark_aiperf.md b/docs/zh/benchmark_aiperf.md index 5f4cb20..0163b37 100644 --- a/docs/zh/benchmark_aiperf.md +++ b/docs/zh/benchmark_aiperf.md @@ -56,10 +56,18 @@ curl --noproxy '*' --fail --silent --show-error \ bash benchmarks/telefuser_aiperf/scripts/run_stream_bench.sh ``` -请求媒体时长为 59.75 秒，但 AIPerf active window 配置为 240 秒，因此命令通常约 4 分钟完成，不要在模型 -停止生成后提前中断。成功输出为 `Stream profile sessions: 1/1 succeeded`，报告写入 -`artifacts/telefuser_aiperf/stream_lingbot_v2_1min/`。模型文件布局、手动选择 Python 环境、历史服务和故障 -排查见上面的 canonical README。 +请求媒体时长为 59.75 秒。240 秒 active window 是超时上限；成功运行会在 target 发出完成状态后退出， +从准入起通常约 66 秒。成功输出为 `Stream profile sessions: 1/1 succeeded`，报告写入 +`artifacts/telefuser_aiperf/stream_lingbot_v2_1min/`。 + +使用相同四卡 workload 测试 SGLang 时，先启动服务，再选择 SGLang 配置： + +```bash +bash benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh + +bash benchmarks/telefuser_aiperf/scripts/run_stream_bench.sh \ + benchmarks/telefuser_aiperf/configs/stream_sglang_lingbot_world_v2_4gpu_1min.json +``` ## 职责与指标语义 @@ -68,6 +76,7 @@ bash benchmarks/telefuser_aiperf/scripts/run_stream_bench.sh | TeleFuser runtime | TeleFuser | 输出同步的 phase、chunk、runtime、cache 和环境原始事实 | | Batch target adapter | AIPerf | 将 `/v1/videos` HTTP 事件转为标准 request 时间线 | | LiveKit 源码 adapter | TeleFuser | 将 room、track、status、metrics 和 control 事件转为 session result | +| SGLang 源码 adapter | TeleFuser | 将 MessagePack frame、chunk timing 和 camera event 转为 session result | | 聚合与历史 | AIPerf | 负责 warmup、percentile、throughput、artifact、GreptimeDB 和展示 | | Contract 与 workload | TeleFuser | 固定 target 能力、输入、设置和可复现启动命令 | @@ -88,35 +97,6 @@ Target 原始事实遵守以下规则： 客户端交付、target pipeline residence、target phase time 和资源利用率保持为不同维度。无法等价的字段保留为 private 或 unavailable，不强行映射为同一指标。 -## LingBot-World v2 一分钟回放实测 - -2026-07-28 使用 4 张 H100 80 GB 验证了 `stream_lingbot_world_v2_1min.json`：DiT 为 BF16，VAE 为 FP32， -启用 SageAttention SM90 和 `torch.compile`，禁用 FSDP，`chunk_size=4`，输出 16 FPS。60 秒请求按完整 latent -chunk 截断为 60 个 chunk、957 帧，对应 59.75 秒媒体时长。LingBot-World v2 使用固定 -`local_attn_size=18`、`sink_size=6` 窗口，因此 240 latent frame 的请求仍只分配 27,144 token KV 容量， -不会改为时长规模的全局 KV cache。 - -| 指标 | 结果 | -|---|---:| -| 成功 session | 1 / 1 | -| Target chunk / 生成帧 | 60 / 957 | -| 客户端收到帧数 | 947 | -| Target 稳态 chunk / 帧 | 59 / 944 | -| 配置的 session runtime | 242.255 s | -| 首帧延迟 | 6,252.773 ms | -| 客户端交付速率（`stream_fps`） | 9.397 FPS | -| 稳态 chunk 加权计算速率 | 4.708 FPS | -| Chunk pipeline residence mean / p99 | 3.399 / 3.901 s | - -Target 实际生成了全部 957 帧，LiveKit 客户端收到 947 帧。AIPerf 在 target 稳态计算聚合中排除了首个 -chunk，剩余 59 个 chunk 共生成 944 帧。242 秒 session runtime 是 240 秒 active window 上限加连接开销， -不是 59.75 秒媒体内容的模型生成耗时。 - -Chunk pipeline residence 从 actor 准入计到输出，包含相互重叠的 encode、DiT、decode 工作，所以可以大于相邻 -交付间隔；报告中的 4.708 稳态 chunk 加权计算 FPS 不能当作客户端推流吞吐。本次 Git 安装的 AIPerf -运行退出码为 0。回放报告位于 -`artifacts/telefuser_aiperf/stream_lingbot_v2_1min/20260728_083948_fc4344ba/stream_report.html`。 - ## 复现要求 每个结果都应保留 TeleFuser commit、AIPerf 包版本、模型 revision、加速器型号/数量、driver、CUDA、 diff --git a/examples/lingbot/README.md b/examples/lingbot/README.md index 52b0158..cb94301 100644 --- a/examples/lingbot/README.md +++ b/examples/lingbot/README.md @@ -22,6 +22,64 @@ Set the model root before running the example: export TF_MODEL_ZOO_PATH=/path/to/model_zoo ``` +## Validated H100 Development Environment + +The four-H100 LingBot-World v2 AIPerf test used the following environment. TeleFuser supports broader +versions through its normal dependency ranges, but performance results in this README should be reproduced with +these versions before attributing a difference to code changes. + +| Component | Validated value | +| --- | --- | +| GPU | 4 x NVIDIA H100 80 GB HBM3 (SM90) | +| NVIDIA driver | `590.48.01` | +| Python | `3.11.13` | +| PyTorch | `2.11.0+cu130` | +| PyTorch CUDA runtime | `13.0` | +| FlashAttention 4 | `flash-attn-4==4.0.0b19` | +| CUTLASS DSL | `nvidia-cutlass-dsl==4.6.0` | +| CUDA Python | `cuda-python==13.3.1` | + +Create an isolated Python 3.11 environment and install the CUDA 13.0 PyTorch build from the wheel index used by +your deployment. Install PyTorch before TeleFuser so optional CUDA packages resolve against the intended ABI: + +```bash +python3.11 -m venv .venv-lingbot +source .venv-lingbot/bin/activate +python -m pip install --upgrade pip setuptools wheel + +# Install torch==2.11.0+cu130 from your CUDA 13.0 PyTorch wheel index first. +python -m pip install -e ".[dev]" +python -m pip install \ + "flash-attn-4[cu13]==4.0.0b19" \ + "nvidia-cutlass-dsl==4.6.0" \ + "cuda-python==13.3.1" +``` + +The `cu13` extra installs FA4's CUDA 13 dependency variant. For a CUDA 12.8 PyTorch environment, install +`flash-attn-4==4.0.0b19` without that extra and use matching CUDA 12.x dependencies; do not mix cu128 and cu130 +interpreters in one distributed run. + +Verify both the package versions and TeleFuser's runtime backend selection before benchmarking: + +```bash +python - <<'PY' +import importlib.metadata as metadata + +import torch + +from telefuser.ops.attention.backends import FLASH_ATTN_4_AVAILABLE + +print("PyTorch:", torch.__version__) +print("PyTorch CUDA:", torch.version.cuda) +print("GPU:", torch.cuda.get_device_name(0)) +print("flash-attn-4:", metadata.version("flash-attn-4")) +print("nvidia-cutlass-dsl:", metadata.version("nvidia-cutlass-dsl")) +print("cuda-python:", metadata.version("cuda-python")) +print("TeleFuser FA4 available:", FLASH_ATTN_4_AVAILABLE) +assert FLASH_ATTN_4_AVAILABLE +PY +``` + ## Feature Support | Feature | Support | @@ -32,7 +90,7 @@ export TF_MODEL_ZOO_PATH=/path/to/model_zoo | Single-GPU inference | ✔️ | | Ulysses Sequence Parallel | ✔️ | | FSDP | Configurable through PPL_CONFIG | -| H100 Sage Attention | ✔️ | +| H100 optimized attention | v2: FA4, then FA3/SageAttention; v1: SageAttention | ## Files @@ -57,8 +115,8 @@ Default configuration: Offline generation and stream-server entry point for camera-controlled v2. The default is 77 frames at 16 FPS: 20 latent frames, exactly five complete chunks of four. With complete chunk streaming, 81 output frames cannot be represented by `chunk_size=4`. -The v2 checkpoint only supports camera control and uses its PPL-configured SageAttention SM90 backend, local attention, sink size, -and timesteps. +The v2 checkpoint only supports camera control. Its H100 example prefers FlashAttention 4, then falls back to FA3 +and SageAttention SM90, while retaining the PPL-configured local attention, sink size, and timesteps. ```bash python examples/lingbot/lingbot_world_v2_image_to_video_h100.py \ @@ -178,7 +236,7 @@ flowchart LR L --> O[StreamingPipelineOrchestrator] O --> E[VAE encode actor] O --> D[DiT actor] - O --> V[VAE decode actor] + O --> V[VAE decode actor or co-located decoder] ``` By default, stream-server calculates `max_sessions_per_worker` after warmup and preallocates fixed DiT KV slots. @@ -197,30 +255,19 @@ fixed placement for the following total GPU counts: | Total GPUs | DiT GPUs | VAE encode GPU | VAE decode GPU | | --- | --- | --- | --- | | 2 | `0-1` | `0` | `1` | -| 4 | `0-3` | `0` | `1` | +| 4 | `0-3` | `0` | `0-3`, co-located with DiT | | 5 | `0-3` | `4` | `4` | | 6 | `0-4` | `5` | `5` | -For other counts, the examples retain the PPL-configured VAE devices and assign all visible GPUs to DiT. Direct -`LingBotWorldFastPipelineConfig` users may set `vae_encode_config`, `vae_decode_config`, and `dit_config` independently. - -### H100 Compile Benchmark - -The v2 example was measured at 480p (832x464 internal size), 77 output frames, five latent chunks of four frames, -BF16 DiT, FP32 VAE, SageAttention SM90, `torch.compile` enabled, and FSDP disabled. Each value is the mean from a -second session after a complete warmup session. Pure DiT measures synchronous `denoise_and_update_cache`; chunk -period is the mean interval between decoded chunk outputs while encode, DiT, and decode overlap. - -| Total H100 GPUs | Pure DiT seconds/chunk | Overlapped chunk period seconds/chunk | -| --- | --- | --- | -| 2 | 1.587 | 2.096 | -| 4 | 0.911 | 1.615 | - -The scheduler does not infer a resource group from overlapping device IDs, so VAE encode, DiT, and VAE decode may -overlap on a shared GPU. +For four GPUs, VAE decode is height-sharded across the same process group as DiT. When the distributed decode and DiT +placements match exactly, the pipeline automatically co-locates them to avoid duplicate CUDA contexts and process +switching. For other counts, the examples retain the PPL-configured VAE devices and assign all visible GPUs to DiT. +Direct `LingBotWorldFastPipelineConfig` users may set `vae_encode_config`, `vae_decode_config`, and `dit_config` +independently; non-matching placements continue to use independent workers. -See the [streaming scheduler guide](../../docs/en/stream_scheduler.md) for -architecture, metric definitions, and lifecycle guarantees. +The scheduler does not infer a resource group from overlapping device IDs. VAE encode remains independent, while an +exactly matching distributed DiT/VAE-decode placement uses the pipeline's explicit co-location path. See the +[streaming scheduler guide](../../docs/en/stream_scheduler.md) for lifecycle guarantees. ### Tested GPU and Duration Limits @@ -238,25 +285,10 @@ The four-GPU 20-second test used FSDP and Ulysses degree 4. Peak memory was appr 41.6 GiB on GPUs 1-3. These are tested values, not universal limits; other resolutions and concurrent GPU users change the available capacity. -LingBot-World v2 instead uses the fixed `local_attn_size=18`, `sink_size=6` sliding window configured by its -example. A four-H100, one-minute AIPerf replay on 2026-07-28 resolved to 60 complete chunks, 957 output frames, and -59.75 seconds of media. Its runtime metadata reported 240 latent frames but a fixed 27,144-token KV capacity. - -| One-minute v2 measurement | Result | -| --- | ---: | -| Successful sessions | 1 / 1 | -| Target chunks / generated frames | 60 / 957 | -| Client frames / steady frames | 946 / 944 | -| Output cadence mean / p50 / p95 | 1.666 / 1.661 / 1.865 s | -| First / middle / last 20-chunk mean | 1.694 / 1.646 / 1.658 s | -| Session runtime | 104.031 s | - -During initial WebRTC track startup, the client received 2 of the first chunk's 13 frames and then received all -`59 * 16 = 944` steady frames. Similar cadence in the first, middle, and last thirds confirms that the fixed attention -window avoided duration-driven degradation in this run. It did not reach real time: each chunk represents 1.0 second -of media, while p95 cadence was 1.865 seconds. See the -[AIPerf benchmark guide](../../docs/en/benchmark_aiperf.md) for the workload, -metric boundary, artifact path, and the observed client cleanup issue. +LingBot-World v2 instead uses a fixed `local_attn_size=18`, `sink_size=6` sliding window, so its cache capacity does +not grow with the one-minute request. The complete four-H100 validation generated 957 frames in 60 chunks at +832x480 without duration-driven cache growth. See the +[benchmark guide](../../docs/en/benchmark_aiperf.md) for delivery results and metric boundaries. ### Camera Controls diff --git a/telefuser/service/livekit/room_client.py b/telefuser/service/livekit/room_client.py index df7554a..06371d0 100644 --- a/telefuser/service/livekit/room_client.py +++ b/telefuser/service/livekit/room_client.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import json from collections.abc import Callable from typing import Any, Protocol @@ -12,13 +13,14 @@ from .token_service import LiveKitDependencyError DataMessageHandler = Callable[[bytes | str | dict[str, Any], str, str], None] -_VIDEO_MAX_BITRATE = 3_000_000 +_VIDEO_MAX_BITRATE = 8_000_000 class RoomClient(Protocol): """Minimal room operations required by a TeleFuser LiveKit worker.""" async def connect(self, url: str, token: str, on_data: DataMessageHandler) -> None: ... + async def wait_for_participant(self, identity: str, *, timeout_s: float) -> None: ... async def publish_video_track(self, name: str, width: int, height: int, *, fps: float = 16.0) -> None: ... async def publish_video_frame(self, frame_rgb: np.ndarray, *, fps: float = 16.0) -> None: ... async def publish_audio_frame(self, pcm: bytes, *, sample_rate: int, channels: int) -> None: ... @@ -60,6 +62,29 @@ def _on_data_received(packet: Any) -> None: await room.connect(url, token) + async def wait_for_participant(self, identity: str, *, timeout_s: float) -> None: + """Wait until a specific remote participant has joined the room.""" + if timeout_s <= 0: + raise ValueError(f"Participant wait timeout must be positive, got {timeout_s}") + room = self._require_room() + if identity in room.remote_participants: + return + + joined = asyncio.Event() + + @room.on("participant_connected") + def _on_participant_connected(participant: Any) -> None: + if getattr(participant, "identity", None) == identity: + joined.set() + + try: + # Cover a participant arriving between the initial check and handler registration. + if identity in room.remote_participants: + return + await asyncio.wait_for(joined.wait(), timeout=timeout_s) + finally: + room.off("participant_connected", _on_participant_connected) + async def publish_video_track(self, name: str, width: int, height: int, *, fps: float = 16.0) -> None: room = self._require_room() rtc = self._require_rtc() diff --git a/telefuser/service/livekit/worker.py b/telefuser/service/livekit/worker.py index f06548c..3dd696d 100644 --- a/telefuser/service/livekit/worker.py +++ b/telefuser/service/livekit/worker.py @@ -22,6 +22,7 @@ from .token_service import LiveKitTokenService _ROOM_DISCONNECT_TIMEOUT_SECONDS = 5.0 +_CONTROLLER_JOIN_TIMEOUT_SECONDS = 60.0 class WorkerEventSink(Protocol): @@ -113,6 +114,10 @@ async def run_session(self, record: SessionRecord) -> None: worker_token, lambda message, topic, identity: self._on_data_message(record, message, topic, identity), ) + await self.room_client.wait_for_participant( + record.controller_identity, + timeout_s=_CONTROLLER_JOIN_TIMEOUT_SECONDS, + ) self.event_sink.on_worker_status(self.worker_id, "starting_pipeline") self.event_sink.on_session_status(record.session_id, "starting_pipeline") @@ -206,6 +211,9 @@ async def _publish_pipeline_chunks( break frames, audio, metadata = split_chunk_media(chunk) + decoded_ready_at = chunk.get("timestamp") + publish_started_at = time.time() + publish_started_monotonic = time.monotonic() chunk_data = chunk.get("data") if isinstance(chunk.get("data"), dict) else chunk fps_value = chunk_data.get("fps", chunk.get("fps", self.config.default_fps)) try: @@ -239,6 +247,14 @@ async def _publish_pipeline_chunks( break if frames: chunk_count += 1 + metadata["transport_measurement"] = { + "decoded_ready_at": decoded_ready_at if isinstance(decoded_ready_at, int | float) else None, + "publish_started_at": publish_started_at, + "publish_finished_at": time.time(), + "publish_seconds": time.monotonic() - publish_started_monotonic, + "frames": len(frames), + "pacing": "realtime", + } if metadata: index = chunk.get("index") if index is None: diff --git a/tests/unit/service/livekit/test_multi_session_worker.py b/tests/unit/service/livekit/test_multi_session_worker.py index 70d39e2..c70f071 100644 --- a/tests/unit/service/livekit/test_multi_session_worker.py +++ b/tests/unit/service/livekit/test_multi_session_worker.py @@ -71,6 +71,9 @@ async def connect(self, url: str, token: str, on_data) -> None: self.on_data = on_data self.connected.set() + async def wait_for_participant(self, identity: str, *, timeout_s: float) -> None: + return None + async def publish_video_track(self, name: str, width: int, height: int, *, fps: float = 16.0) -> None: return None diff --git a/tests/unit/service/livekit/test_room_client.py b/tests/unit/service/livekit/test_room_client.py index 2dd98a7..d8d66b0 100644 --- a/tests/unit/service/livekit/test_room_client.py +++ b/tests/unit/service/livekit/test_room_client.py @@ -84,6 +84,7 @@ class FakeRoom: def __init__(self) -> None: self.local_participant = FakeLocalParticipant() self.handlers = {} + self.remote_participants = {} def on(self, event: str): def _decorator(fn): @@ -92,6 +93,10 @@ def _decorator(fn): return _decorator + def off(self, event: str, fn) -> None: + if self.handlers.get(event) is fn: + self.handlers.pop(event) + async def connect(self, url: str, token: str) -> None: captured["connect"] = (url, token) @@ -127,6 +132,13 @@ async def _run() -> None: packet = types.SimpleNamespace(data=b"{}", topic="tf.control", participant=participant) fake_room.handlers["data_received"](packet) + controller = types.SimpleNamespace(identity="controller") + wait_task = asyncio.create_task(client.wait_for_participant("controller", timeout_s=1.0)) + await asyncio.sleep(0) + fake_room.remote_participants["controller"] = controller + fake_room.handlers["participant_connected"](controller) + await wait_task + frame = np.zeros((2, 3, 3), dtype=np.uint8) await client.publish_video_frame(frame, fps=16) _video_track, video_options = captured["published_track"] @@ -145,7 +157,7 @@ async def _run() -> None: assert video_options.simulcast is False assert video_options.video_codec == "VP8" assert video_options.video_encoding.max_framerate == 16 - assert video_options.video_encoding.max_bitrate == 3_000_000 + assert video_options.video_encoding.max_bitrate == 8_000_000 assert captured["frame"]["buffer_type"] == "RGB24" assert captured["audio_source"] == (48_000, 1) assert captured["audio_frame"] == (pcm, 48_000, 1, 960) diff --git a/tests/unit/service/livekit/test_worker.py b/tests/unit/service/livekit/test_worker.py index acaea6d..82a8ca7 100644 --- a/tests/unit/service/livekit/test_worker.py +++ b/tests/unit/service/livekit/test_worker.py @@ -76,12 +76,19 @@ def __init__(self) -> None: self.statuses: list[dict] = [] self.disconnected = False self.disconnect_gate: asyncio.Event | None = None + self.participant_gate = asyncio.Event() + self.participant_gate.set() + self.waited_for_participants: list[tuple[str, float]] = [] async def connect(self, url: str, token: str, on_data) -> None: self.connect_args = (url, token) self.on_data = on_data self.connected.set() + async def wait_for_participant(self, identity: str, *, timeout_s: float) -> None: + self.waited_for_participants.append((identity, timeout_s)) + await asyncio.wait_for(self.participant_gate.wait(), timeout=timeout_s) + async def publish_video_track(self, name: str, width: int, height: int, *, fps: float = 16.0) -> None: return None @@ -145,6 +152,7 @@ def _native_chunk() -> dict: "type": "chunk", "index": 1, "fps": 16, + "timestamp": 123.0, "frames": [Image.new("RGB", (8, 8), color=(1, 2, 3))], "stream_progress": {"completed_chunks": 2}, } @@ -215,11 +223,21 @@ async def _run() -> None: assert adapter.pushed == [("pipeline-session-1", {"type": "control", "event": "press", "key": "ArrowUp"})] assert adapter.closed == ["pipeline-session-1"] assert room.connect_args == ("wss://livekit.example", "worker:telefuser-worker-0:room-1") + assert room.waited_for_participants == [("controller", worker_module._CONTROLLER_JOIN_TIMEOUT_SECONDS)] assert len(room.video_frames) == 2 assert room.video_frame_fps == [16.0, 16.0] assert room.audio_frames == [(np.zeros(960, dtype=np.int16).tobytes(), 48_000, 1)] assert any(status.get("data", {}).get("frames") == 13 for status in room.statuses) assert any(status.get("data", {}).get("stream_progress") == {"completed_chunks": 2} for status in room.statuses) + transport = next( + status["data"]["transport_measurement"] + for status in room.statuses + if status.get("data", {}).get("index") == 1 + ) + assert transport["decoded_ready_at"] == 123.0 + assert transport["pacing"] == "realtime" + assert transport["frames"] == 1 + assert transport["publish_started_at"] <= transport["publish_finished_at"] assert room.statuses[-1]["type"] == "done" assert room.disconnected is True assert sink.pipeline_sessions == [("session-1", "pipeline-session-1")] @@ -229,6 +247,47 @@ async def _run() -> None: asyncio.run(_run()) +def test_livekit_worker_waits_for_controller_before_creating_pipeline() -> None: + async def _run() -> None: + adapter = FakePipelineAdapter() + room = FakeRoomClient() + room.participant_gate.clear() + worker = LiveKitWorker( + worker_id="worker-0", + config=LiveKitServeConfig( + livekit_url="wss://livekit.example", + livekit_api_key="key", + livekit_api_secret="secret", + ), + pipeline_file="pipeline.py", + token_service=FakeTokenService(), + pipeline_adapter=adapter, + room_client=room, + ) + record = SessionRecord( + session_id="session-1", + room_name="room-1", + controller_identity="controller", + status="assigned", + worker_id="worker-0", + config={"session_id": "session-1"}, + created_at=0, + updated_at=0, + ) + + task = asyncio.create_task(worker.run_session(record)) + await room.connected.wait() + await asyncio.sleep(0) + assert adapter.created_config is None + + room.participant_gate.set() + await adapter.created.wait() + await adapter.output_queue.put(None) + await task + + asyncio.run(_run()) + + def test_livekit_worker_runs_server_push_pipeline() -> None: async def _run() -> None: adapter = FakePipelineAdapter(stream_mode=STREAM_MODE_SERVER_PUSH) diff --git a/tools/validation/benchmark_lingbot_world_v2_direct.py b/tools/validation/benchmark_lingbot_world_v2_direct.py new file mode 100644 index 0000000..7710539 --- /dev/null +++ b/tools/validation/benchmark_lingbot_world_v2_direct.py @@ -0,0 +1,174 @@ +"""Benchmark LingBot-World v2 through the direct pipeline-service output path.""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import platform +import statistics +import sys +import time +from pathlib import Path +from typing import Any + +import torch + +from telefuser.service.livekit.pipeline_adapter import LiveKitPipelineAdapter +from telefuser.service.security.security_validator import SecurityLevel + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--pipeline", required=True) + parser.add_argument("--image", required=True) + parser.add_argument("--control-trace", required=True) + parser.add_argument("--output", required=True) + parser.add_argument("--gpu-num", type=int, default=4) + parser.add_argument("--prompt", default="walk forward through the scene") + parser.add_argument("--frame-num", type=int, default=957) + parser.add_argument("--fps", type=int, default=16) + parser.add_argument("--chunk-size", type=int, default=4) + return parser.parse_args() + + +def _summary(values: list[float]) -> dict[str, float | int]: + ordered = sorted(values) + p90_index = int(0.9 * (len(ordered) - 1)) + return { + "count": len(values), + "mean": statistics.mean(values), + "p50": statistics.median(values), + "p90": ordered[p90_index], + "max": ordered[-1], + } + + +async def _send_controls( + adapter: LiveKitPipelineAdapter, + session_id: str, + events: list[dict[str, Any]], + started_at: float, +) -> None: + for event in events: + delay = started_at + float(event["delay_s"]) - time.monotonic() + if delay > 0: + await asyncio.sleep(delay) + adapter.push_chunk(session_id, dict(event["message"])) + + +async def _run(args: argparse.Namespace) -> dict[str, Any]: + trace = json.loads(Path(args.control_trace).read_text()) + events = trace["events"] + adapter = LiveKitPipelineAdapter(security_level=SecurityLevel.NONE) + adapter.start(args.pipeline, skip_validation=True, gpu_num=args.gpu_num) + capacity = adapter.configure_session_capacity(2) + session_id = adapter.create_session( + { + "prompt": args.prompt, + "image_path": str(Path(args.image).resolve()), + "fps": args.fps, + "chunk_size": args.chunk_size, + "frame_num": args.frame_num, + "max_duration_seconds": 60.0, + "sample_shift": 10.0, + "control_mode": "cam", + "show_control_hud": False, + "benchmark_metrics": True, + } + ) + started_at = time.monotonic() + sender = asyncio.create_task(_send_controls(adapter, session_id, events, started_at)) + frames = 0 + chunk_profiles: list[dict[str, Any]] = [] + runtime: dict[str, Any] | None = None + first_preview_at: float | None = None + first_generated_frame_at: float | None = None + try: + async for payload in adapter.pull_chunks(session_id): + payload_type = payload.get("type") + if payload_type in {"preview", "chunk"}: + payload_frames = payload.get("frames", []) + if payload_type == "chunk": + frames += len(payload_frames) + if payload_frames and first_generated_frame_at is None: + first_generated_frame_at = time.monotonic() + elif payload_frames and first_preview_at is None: + first_preview_at = time.monotonic() + if payload_type != "status": + continue + if payload.get("stage") == "runtime_ready": + runtime = payload.get("runtime") + measurement = payload.get("measurement") + if isinstance(measurement, dict) and "index" in measurement: + chunk_profiles.append(measurement) + await sender + finally: + if not sender.done(): + sender.cancel() + await adapter.aclose() + + elapsed = time.monotonic() - started_at + steady = [profile for profile in chunk_profiles if int(profile["index"]) > 0] + compute_seconds = [float(profile["compute_seconds"]) for profile in steady] + phase_names = sorted( + { + name + for profile in steady + for name, value in profile.get("phases", {}).items() + if isinstance(value, (int, float)) and not isinstance(value, bool) + } + ) + return { + "environment": { + "sys_executable": sys.executable, + "python": platform.python_version(), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + }, + "request": { + "pipeline": str(Path(args.pipeline).resolve()), + "image": str(Path(args.image).resolve()), + "control_trace": str(Path(args.control_trace).resolve()), + "prompt": args.prompt, + "frame_num": args.frame_num, + "fps": args.fps, + "chunk_size": args.chunk_size, + "gpu_num": args.gpu_num, + }, + "transport": "direct pipeline service; no LiveKit room, pacing, codec, or client", + "capacity": capacity, + "runtime": runtime, + "result": { + "frames": frames, + "chunks": len(chunk_profiles), + "steady_chunks": len(steady), + "elapsed_seconds": elapsed, + "first_preview_seconds": None if first_preview_at is None else first_preview_at - started_at, + "first_generated_frame_seconds": ( + None if first_generated_frame_at is None else first_generated_frame_at - started_at + ), + "steady_compute_seconds": sum(compute_seconds), + "steady_compute_fps": sum(float(profile["frames"]) for profile in steady) / sum(compute_seconds), + "chunk_compute_seconds": _summary(compute_seconds), + "phases": { + name: _summary([float(profile["phases"][name]) for profile in steady if name in profile["phases"]]) + for name in phase_names + }, + }, + "chunk_profiles": chunk_profiles, + } + + +def main() -> None: + args = _parse_args() + result = asyncio.run(_run(args)) + output = Path(args.output) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n") + print(json.dumps(result["result"], indent=2, sort_keys=True)) + print(f"Artifact: {output}") + + +if __name__ == "__main__": + main() From 8cfd68b267414a222d13736dbbe37737a0c46bfa Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Sun, 2 Aug 2026 10:08:58 +0000 Subject: [PATCH 05/11] feat(worker): add direct tensor channels Transport tensor outputs between independent worker groups with metadata references while CUDA IPC or shared memory carries the storage directly. Move final device placement into consumer ranks and allow method-scoped tensor transport without routing payloads through the parent process.\n\nDocument the channel lifecycle and cover FIFO, cancellation, CPU sharing, CUDA peer copies, and ParallelWorker bindings.\n\nVerification: focused worker unit tests and the two-GPU tensor-channel integration test passed; ruff and git diff --check passed. --- docs/en/parallel.md | 28 +++ docs/zh/parallel.md | 27 +++ telefuser/worker/__init__.py | 3 + telefuser/worker/parallel_worker.py | 47 +++- telefuser/worker/tensor_channel.py | 183 +++++++++++++++ .../integration/test_worker_tensor_channel.py | 64 ++++++ tests/unit/worker/test_parallel_worker.py | 45 ++++ tests/unit/worker/test_tensor_channel.py | 213 ++++++++++++++++++ 8 files changed, 603 insertions(+), 7 deletions(-) create mode 100644 telefuser/worker/tensor_channel.py create mode 100644 tests/integration/test_worker_tensor_channel.py create mode 100644 tests/unit/worker/test_tensor_channel.py diff --git a/docs/en/parallel.md b/docs/en/parallel.md index a846346..4129711 100644 --- a/docs/en/parallel.md +++ b/docs/en/parallel.md @@ -159,6 +159,34 @@ out = out_wait() Split model layers across multiple GPUs for large model inference. +### Cross-worker tensor channels + +Independent `ParallelWorker` groups can connect adjacent stages with a `WorkerTensorChannel`. The producer sends +tensor storage directly to every consumer rank through multiprocessing shared memory or CUDA IPC. The parent process +receives only a `WorkerTensorRef` containing the channel, transfer, shape, dtype, and source-device metadata. The +consumer resolves that reference on its own device before invoking the unchanged stage method. + +```python +from telefuser.worker import ParallelWorker, WorkerTensorChannel + +latent_channel = WorkerTensorChannel(consumer_world_size=vae_parallel_config.world_size) +denoise_worker = ParallelWorker( + denoise_stage, + tensor_output_channel=latent_channel, + tensor_output_methods=("denoise",), +) +vae_worker = ParallelWorker(vae_stage, tensor_input_channels=(latent_channel,)) +``` + +This is a point-to-point, single-producer/single-consumer-group path. Enable it only for outputs whose complete +consumer set is the connected worker group. Calls that need to inspect a tensor in the parent may pass +`_tensor_transport=False`. Start both workers before submitting work, stop both workers before closing the channel, +preserve producer order at the consumer, and treat transported tensors as immutable until the consumer finishes. +The receiver discards older FIFO entries when the scheduler cancels an artifact before consumption. + +Regular worker dispatch also sends shared-memory or CUDA IPC handles to each rank and lets the receiving rank perform +the final device placement. The parent does not allocate a temporary copy on every target GPU. + ### Principle ``` diff --git a/docs/zh/parallel.md b/docs/zh/parallel.md index 1ad2ee7..b3e7a20 100644 --- a/docs/zh/parallel.md +++ b/docs/zh/parallel.md @@ -159,6 +159,33 @@ out = out_wait() 将模型层分割到多个 GPU，实现大模型推理。 +### 跨 Worker Tensor 通道 + +相邻 stage 如果属于不同 `ParallelWorker` group，可以使用 `WorkerTensorChannel` 连接。Producer 通过 +multiprocessing shared memory 或 CUDA IPC，把 tensor storage 直接发送给每个 consumer rank。主进程只接收 +`WorkerTensorRef`，其中包含 channel、transfer、shape、dtype、字节数和源设备等元数据。Consumer 在调用原有 +stage 方法之前，负责在自己的设备上解析引用；stage 的函数签名无需改变。 + +```python +from telefuser.worker import ParallelWorker, WorkerTensorChannel + +latent_channel = WorkerTensorChannel(consumer_world_size=vae_parallel_config.world_size) +denoise_worker = ParallelWorker( + denoise_stage, + tensor_output_channel=latent_channel, + tensor_output_methods=("denoise",), +) +vae_worker = ParallelWorker(vae_stage, tensor_input_channels=(latent_channel,)) +``` + +该路径是单 producer、单 consumer group 的点对点 FIFO。只有 tensor 的完整 consumer 集合就是所连接的 worker +group 时才能启用。需要在主进程读取 tensor 的调用可以传入 `_tensor_transport=False`。Consumer 必须保持 +producer 顺序，并在两个 worker 都停止后再关闭 channel。Scheduler 取消尚未消费的 artifact 时，receiver 会 +丢弃更早的 FIFO entry，避免污染后续传输。 + +常规 worker 派发同样只向各 rank 发送 shared-memory 或 CUDA IPC handle，最终 device placement 由接收 rank +完成；主进程不再为每张目标 GPU 分配临时副本。 + ### 原理 ``` diff --git a/telefuser/worker/__init__.py b/telefuser/worker/__init__.py index 72f0fb2..ccbc685 100644 --- a/telefuser/worker/__init__.py +++ b/telefuser/worker/__init__.py @@ -8,9 +8,12 @@ from .parallel_worker import ParallelWorker from .ray_worker import RayWorker, create_ray_worker +from .tensor_channel import WorkerTensorChannel, WorkerTensorRef __all__ = [ "ParallelWorker", "RayWorker", + "WorkerTensorChannel", + "WorkerTensorRef", "create_ray_worker", ] diff --git a/telefuser/worker/parallel_worker.py b/telefuser/worker/parallel_worker.py index b19b840..394858b 100644 --- a/telefuser/worker/parallel_worker.py +++ b/telefuser/worker/parallel_worker.py @@ -10,7 +10,7 @@ import os import threading import time -from collections.abc import Callable +from collections.abc import Callable, Collection from datetime import timedelta from multiprocessing.queues import SimpleQueue from queue import Empty @@ -26,6 +26,8 @@ from telefuser.utils.logging import logger from telefuser.utils.system import PortAllocator +from .tensor_channel import WorkerTensorChannel + if TYPE_CHECKING: from telefuser.metrics import StageMetricContext @@ -42,7 +44,7 @@ def to_device(data: Any, device: str | torch.device) -> Any: if device == "cpu" and not data.is_shared(): data.share_memory_() return data - tensor = data.clone().to(device) + tensor = data.to(device) if device == "cpu" and not tensor.is_shared(): tensor.share_memory_() return tensor @@ -57,6 +59,8 @@ def _worker_loop( queue_out: SimpleQueue, stage: BaseStage, master_port: int, + tensor_output_channel: WorkerTensorChannel | None = None, + tensor_input_channels: tuple[WorkerTensorChannel, ...] = (), ) -> None: """Worker process main loop. @@ -108,13 +112,22 @@ def _worker_loop( while True: data = queue_in[rank].get() - name, args, kwargs = data + if len(data) == 3: + name, args, kwargs = data + transport_output = False + else: + name, args, kwargs, transport_output = data del data if name == "exit": logger.info(f"parallel worker {stage.name} on rank {rank} exits") break if not hasattr(stage, name): raise AttributeError(f'{stage.__class__.__name__} has no attribute "{name}"') + stage_inputs = (args, kwargs) + for channel in tensor_input_channels: + if channel.contains_ref(stage_inputs): + stage_inputs = channel.receive(stage_inputs, rank=rank, device=device) + args, kwargs = stage_inputs kwargs = to_device(kwargs, device) args = to_device(args, device) with torch.no_grad(): @@ -124,6 +137,8 @@ def _worker_loop( del kwargs, args if getattr(stage, "empty_cache_after_call", True): current_platform.empty_cache() + if transport_output and tensor_output_channel is not None and (world_size == 1 or rank == 0): + y = tensor_output_channel.send(y) # Always output results when world_size=1 if world_size == 1 or rank == 0: queue_out.put(y) @@ -155,6 +170,10 @@ class ParallelWorker: def __init__( self, stage: BaseStage, + *, + tensor_output_channel: WorkerTensorChannel | None = None, + tensor_output_methods: Collection[str] = (), + tensor_input_channels: Collection[WorkerTensorChannel] = (), ) -> None: parallel_config = stage.model_runtime_config.parallel_config parallel_config.validate() @@ -173,6 +192,17 @@ def __init__( self._failed = False self._closed = False self._failure_reason: str | None = None + self.tensor_output_channel = tensor_output_channel + self.tensor_output_methods = frozenset(tensor_output_methods) + self.tensor_input_channels = tuple(tensor_input_channels) + if self.tensor_output_channel is None and self.tensor_output_methods: + raise ValueError("tensor_output_methods require a tensor_output_channel") + if self.tensor_output_channel is not None: + if not self.tensor_output_methods: + raise ValueError("tensor_output_channel requires at least one tensor_output_method") + self.tensor_output_channel.bind_producer() + for channel in self.tensor_input_channels: + channel.bind_consumer(self.world_size) # Use spawn to start processes regardless of world_size current_method = mp.get_start_method(allow_none=True) @@ -201,6 +231,8 @@ def __init__( self.queue_out, stage, master_port, + self.tensor_output_channel, + self.tensor_input_channels, ), nprocs=self.world_size, join=False, @@ -297,15 +329,15 @@ def put_data(self, data: Any) -> None: self._ensure_usable() if self.queue_with_cpu: data = to_device(data, "cpu") - for i, q in enumerate(self.queue_in): - data = to_device(data, device=f"{self.device}:{self.device_ids[i]}") + for q in self.queue_in: q.put(data) def __call__(self, *args: Any, **kwargs: Any) -> Any | Callable[[], Any]: """Submit __call__ task to all workers.""" self._ensure_usable() sync = kwargs.pop("sync", False) - data = ["__call__", args, kwargs] + transport_output = kwargs.pop("_tensor_transport", "__call__" in self.tensor_output_methods) + data = ["__call__", args, kwargs, transport_output] self.put_data(data) def wait() -> Any: @@ -322,7 +354,8 @@ def __getattr__(self, name: str) -> Callable[..., Any]: def wrapped_func(*args: Any, **kwargs: Any) -> Any | Callable[[], Any]: self._ensure_usable() sync = kwargs.pop("sync", False) - data = [name, args, kwargs] + transport_output = kwargs.pop("_tensor_transport", name in self.tensor_output_methods) + data = [name, args, kwargs, transport_output] self.put_data(data) hook = self._metrics_hook diff --git a/telefuser/worker/tensor_channel.py b/telefuser/worker/tensor_channel.py new file mode 100644 index 0000000..6046a2a --- /dev/null +++ b/telefuser/worker/tensor_channel.py @@ -0,0 +1,183 @@ +"""Direct tensor transport between independently spawned worker groups.""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass +from multiprocessing.queues import SimpleQueue +from typing import Any + +import torch +import torch.multiprocessing as mp + + +@dataclass(frozen=True) +class WorkerTensorRef: + """Metadata-only reference to a tensor held by a worker-to-worker channel.""" + + channel_id: str + transfer_id: int + tensor_index: int + shape: tuple[int, ...] + dtype: str + source_device: str + nbytes: int + + +class WorkerTensorChannel: + """Point-to-point tensor path that bypasses the parent process. + + The producer places tensors directly onto one queue per consumer rank. CUDA + tensors travel as CUDA IPC handles; CPU tensors use multiprocessing shared + memory. The parent process receives only :class:`WorkerTensorRef` objects. + """ + + def __init__(self, consumer_world_size: int, *, timeout: int = 600) -> None: + if consumer_world_size < 1: + raise ValueError("consumer_world_size must be at least one") + if timeout < 1: + raise ValueError("timeout must be at least one second") + spawn_ctx = mp.get_context("spawn") + self.channel_id = uuid.uuid4().hex + self.consumer_world_size = consumer_world_size + self.timeout = timeout + self._queues: tuple[SimpleQueue, ...] = tuple(spawn_ctx.SimpleQueue() for _ in range(consumer_world_size)) + self._next_transfer_id = 0 + self._producer_bound = False + self._consumer_bound = False + self._closed = False + + def bind_producer(self) -> None: + """Bind exactly one producer worker group.""" + if self._producer_bound: + raise ValueError(f"Tensor channel {self.channel_id} already has a producer") + if self._closed: + raise RuntimeError(f"Tensor channel {self.channel_id} is closed") + self._producer_bound = True + + def bind_consumer(self, world_size: int) -> None: + """Bind exactly one consumer worker group with a matching rank count.""" + if world_size != self.consumer_world_size: + raise ValueError( + f"Tensor channel {self.channel_id} expects {self.consumer_world_size} consumer ranks, got {world_size}" + ) + if self._consumer_bound: + raise ValueError(f"Tensor channel {self.channel_id} already has a consumer") + if self._closed: + raise RuntimeError(f"Tensor channel {self.channel_id} is closed") + self._consumer_bound = True + + def send(self, value: Any) -> Any: + """Send every tensor leaf to all consumer ranks and return metadata refs.""" + transfer_id = self._next_transfer_id + self._next_transfer_id += 1 + tensor_index = 0 + sent: dict[int, WorkerTensorRef] = {} + + def replace(item: Any) -> Any: + nonlocal tensor_index + if isinstance(item, torch.Tensor): + existing = sent.get(id(item)) + if existing is not None: + return existing + ref = WorkerTensorRef( + channel_id=self.channel_id, + transfer_id=transfer_id, + tensor_index=tensor_index, + shape=tuple(item.shape), + dtype=str(item.dtype), + source_device=str(item.device), + nbytes=item.numel() * item.element_size(), + ) + tensor_index += 1 + sent[id(item)] = ref + for queue in self._queues: + queue.put((ref, item)) + return ref + if isinstance(item, dict): + return {key: replace(child) for key, child in item.items()} + if isinstance(item, tuple): + return tuple(replace(child) for child in item) + if isinstance(item, list): + return [replace(child) for child in item] + return item + + return replace(value) + + def receive(self, value: Any, *, rank: int, device: str | torch.device) -> Any: + """Resolve tensor refs for one consumer rank onto its local device.""" + if not 0 <= rank < self.consumer_world_size: + raise ValueError(f"Consumer rank {rank} is outside [0, {self.consumer_world_size})") + queue = self._queues[rank] + resolved: dict[WorkerTensorRef, torch.Tensor] = {} + + def receive_tensor(expected: WorkerTensorRef) -> torch.Tensor: + while True: + if not queue._reader.poll(self.timeout): + raise TimeoutError( + f"Tensor channel {self.channel_id} timed out receiving transfer {expected.transfer_id}" + ) + received_ref, tensor = queue.get() + received_key = (received_ref.transfer_id, received_ref.tensor_index) + expected_key = (expected.transfer_id, expected.tensor_index) + if received_key < expected_key: + # The parent dropped this earlier artifact, normally after + # cancellation. Releasing it here keeps the FIFO usable. + del tensor + continue + if received_ref != expected: + raise RuntimeError( + f"Tensor channel {self.channel_id} expected {expected}, received {received_ref}; " + "consumer calls must preserve producer order" + ) + return tensor + + def replace(item: Any) -> Any: + if isinstance(item, WorkerTensorRef): + if item.channel_id != self.channel_id: + return item + cached = resolved.get(item) + if cached is not None: + return cached + tensor = receive_tensor(item) + if ( + tuple(tensor.shape) != item.shape + or str(tensor.dtype) != item.dtype + or str(tensor.device) != item.source_device + or tensor.numel() * tensor.element_size() != item.nbytes + ): + raise RuntimeError( + f"Tensor channel {self.channel_id} received incompatible tensor metadata for {item}" + ) + target = torch.device(device) + if tensor.device != target: + tensor = tensor.to(target, non_blocking=True) + resolved[item] = tensor + return tensor + if isinstance(item, dict): + return {key: replace(child) for key, child in item.items()} + if isinstance(item, tuple): + return tuple(replace(child) for child in item) + if isinstance(item, list): + return [replace(child) for child in item] + return item + + return replace(value) + + def contains_ref(self, value: Any) -> bool: + """Return whether a nested value contains a ref owned by this channel.""" + if isinstance(value, WorkerTensorRef): + return value.channel_id == self.channel_id + if isinstance(value, dict): + return any(self.contains_ref(child) for child in value.values()) + if isinstance(value, tuple | list): + return any(self.contains_ref(child) for child in value) + return False + + def close(self) -> None: + """Close parent-owned queue handles after both worker groups stop.""" + if self._closed: + return + self._closed = True + for queue in self._queues: + queue.close() diff --git a/tests/integration/test_worker_tensor_channel.py b/tests/integration/test_worker_tensor_channel.py new file mode 100644 index 0000000..caee279 --- /dev/null +++ b/tests/integration/test_worker_tensor_channel.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import pytest +import torch +import torch.distributed as dist + +from telefuser.core.base_stage import BaseStage +from telefuser.core.config import ModelRuntimeConfig, ParallelConfig +from telefuser.worker import ParallelWorker, WorkerTensorChannel, WorkerTensorRef + + +class _DistributedProducerStage(BaseStage): + def __init__(self) -> None: + super().__init__( + "distributed-producer", + ModelRuntimeConfig( + device_type="cuda", + device_id=0, + parallel_config=ParallelConfig(device_ids=[0, 1], sp_ulysses_degree=2), + ), + ) + + def reduce(self, tensor: torch.Tensor) -> torch.Tensor: + result = tensor + dist.get_rank() + dist.all_reduce(result) + return result + + +class _GPUConsumerStage(BaseStage): + def __init__(self) -> None: + super().__init__( + "gpu-consumer", + ModelRuntimeConfig( + device_type="cuda", + device_id=2, + parallel_config=ParallelConfig(device_ids=[2]), + ), + ) + + def consume(self, tensor: torch.Tensor) -> tuple[str, torch.Tensor]: + return str(tensor.device), tensor.cpu() + + +@pytest.mark.distributed +@pytest.mark.skipif(torch.cuda.device_count() < 3, reason="requires three CUDA devices") +def test_distributed_worker_to_worker_tensor_path_bypasses_parent() -> None: + channel = WorkerTensorChannel(consumer_world_size=1, timeout=30) + producer = ParallelWorker( + _DistributedProducerStage(), + tensor_output_channel=channel, + tensor_output_methods=("reduce",), + ) + consumer = ParallelWorker(_GPUConsumerStage(), tensor_input_channels=(channel,)) + try: + ref = producer.reduce(torch.arange(4, dtype=torch.float32), sync=True) + assert isinstance(ref, WorkerTensorRef) + assert ref.source_device == "cuda:0" + device, result = consumer.consume(ref, sync=True) + assert device == "cuda:2" + torch.testing.assert_close(result, 2 * torch.arange(4, dtype=torch.float32) + 1) + finally: + producer.close() + consumer.close() + channel.close() diff --git a/tests/unit/worker/test_parallel_worker.py b/tests/unit/worker/test_parallel_worker.py index f3a0596..43cc6d4 100644 --- a/tests/unit/worker/test_parallel_worker.py +++ b/tests/unit/worker/test_parallel_worker.py @@ -142,6 +142,51 @@ def test_put_data_with_queue_cpu(self, mock_port_allocator, mock_mp): mock_queue.put.assert_called_once() + @patch("telefuser.worker.parallel_worker.to_device") + def test_put_data_broadcasts_handles_without_parent_device_copies(self, mock_to_device): + """Target-device copies belong to ranks, not the parent dispatch loop.""" + from telefuser.worker.parallel_worker import ParallelWorker + + worker = ParallelWorker.__new__(ParallelWorker) + worker._closed = False + worker._failed = False + worker.queue_with_cpu = False + worker.queue_in = [MagicMock(), MagicMock()] + payload = ["method", (torch.ones(1),), {}, False] + + worker.put_data(payload) + + mock_to_device.assert_not_called() + for queue in worker.queue_in: + queue.put.assert_called_once_with(payload) + + @patch("telefuser.worker.parallel_worker.mp") + @patch("telefuser.worker.parallel_worker.PortAllocator") + def test_initialization_binds_direct_tensor_channels(self, mock_port_allocator, mock_mp): + from telefuser.worker.parallel_worker import ParallelWorker + + mock_port_allocator.return_value.get_free_port_in_interval.return_value = 12345 + mock_mp.get_start_method.return_value = "spawn" + mock_mp.get_context.return_value = MagicMock() + output_channel = MagicMock() + input_channel = MagicMock() + mock_stage = MagicMock() + mock_stage.name = "TestStage" + mock_stage.model_runtime_config.parallel_config.world_size = 2 + mock_stage.model_runtime_config.parallel_config.device_ids = [0, 1] + mock_stage.model_runtime_config.parallel_config.queue_with_cpu = False + mock_stage.model_runtime_config.parallel_config.timeout = 600 + + ParallelWorker( + mock_stage, + tensor_output_channel=output_channel, + tensor_output_methods=("forward",), + tensor_input_channels=(input_channel,), + ) + + output_channel.bind_producer.assert_called_once_with() + input_channel.bind_consumer.assert_called_once_with(2) + @patch("telefuser.worker.parallel_worker.mp") @patch("telefuser.worker.parallel_worker.PortAllocator") def test_call_with_sync(self, mock_port_allocator, mock_mp): diff --git a/tests/unit/worker/test_tensor_channel.py b/tests/unit/worker/test_tensor_channel.py new file mode 100644 index 0000000..71fd17f --- /dev/null +++ b/tests/unit/worker/test_tensor_channel.py @@ -0,0 +1,213 @@ +from __future__ import annotations + +import pytest +import torch +import torch.multiprocessing as mp + +from telefuser.core.base_stage import BaseStage +from telefuser.core.config import ModelRuntimeConfig, ParallelConfig +from telefuser.worker.parallel_worker import ParallelWorker +from telefuser.worker.tensor_channel import WorkerTensorChannel, WorkerTensorRef + + +class _TensorProducerStage(BaseStage): + def __init__(self) -> None: + super().__init__( + "tensor-producer", + ModelRuntimeConfig(device_type="cpu", parallel_config=ParallelConfig(device_ids=[0])), + ) + + def produce(self) -> tuple[torch.Tensor, dict[str, float]]: + return torch.arange(4, dtype=torch.float32), {"seconds": 0.25} + + +class _TensorConsumerStage(BaseStage): + def __init__(self) -> None: + super().__init__( + "tensor-consumer", + ModelRuntimeConfig(device_type="cpu", parallel_config=ParallelConfig(device_ids=[0])), + ) + + def consume(self, tensor: torch.Tensor) -> torch.Tensor: + return tensor + 1 + + +def _spawn_send(channel, metadata_queue, release_event) -> None: + metadata_queue.put(channel.send(torch.arange(4, dtype=torch.float32))) + release_event.wait(timeout=10) + + +def _spawn_receive(channel, metadata_queue, result_queue, release_event) -> None: + ref = metadata_queue.get() + tensor = channel.receive(ref, rank=0, device="cpu") + result_queue.put(tensor.tolist()) + release_event.set() + + +def _spawn_send_cuda(channel, metadata_queue, release_event) -> None: + torch.cuda.set_device(0) + tensor = torch.arange(4, dtype=torch.float32, device="cuda:0") + metadata_queue.put(channel.send(tensor)) + release_event.wait(timeout=30) + + +def _spawn_receive_cuda(channel, metadata_queue, result_queue, release_event) -> None: + torch.cuda.set_device(1) + ref = metadata_queue.get() + tensor = channel.receive(ref, rank=0, device="cuda:1") + torch.cuda.synchronize(1) + result_queue.put((str(tensor.device), tensor.cpu().tolist())) + release_event.set() + + +def test_tensor_channel_keeps_parent_artifact_metadata_only_and_fans_out() -> None: + channel = WorkerTensorChannel(consumer_world_size=2, timeout=1) + channel.bind_producer() + channel.bind_consumer(2) + source = torch.arange(6, dtype=torch.float32).reshape(2, 3) + + try: + artifact = channel.send({"latent": source, "profile": {"seconds": 0.1}}) + + assert isinstance(artifact["latent"], WorkerTensorRef) + assert artifact["latent"].shape == (2, 3) + assert artifact["latent"].nbytes == source.numel() * source.element_size() + assert artifact["profile"] == {"seconds": 0.1} + for rank in range(2): + resolved = channel.receive(artifact, rank=rank, device="cpu") + torch.testing.assert_close(resolved["latent"], source) + assert resolved["profile"] == {"seconds": 0.1} + finally: + channel.close() + + +def test_tensor_channel_preserves_nested_container_types() -> None: + channel = WorkerTensorChannel(consumer_world_size=1, timeout=1) + try: + artifact = channel.send((torch.ones(1), [torch.zeros(1)])) + resolved = channel.receive(artifact, rank=0, device="cpu") + finally: + channel.close() + + assert isinstance(artifact, tuple) + assert isinstance(artifact[1], list) + assert isinstance(resolved, tuple) + assert isinstance(resolved[1], list) + + +def test_tensor_channel_sends_duplicate_tensor_leaf_once() -> None: + channel = WorkerTensorChannel(consumer_world_size=1, timeout=1) + source = torch.ones(2) + try: + artifact = channel.send((source, {"same": source})) + assert artifact[0] == artifact[1]["same"] + resolved = channel.receive(artifact, rank=0, device="cpu") + finally: + channel.close() + + assert resolved[0] is resolved[1]["same"] + + +def test_tensor_channel_discards_cancelled_earlier_transfer() -> None: + channel = WorkerTensorChannel(consumer_world_size=1, timeout=1) + try: + channel.send(torch.zeros(1)) + current = channel.send(torch.ones(1)) + resolved = channel.receive(current, rank=0, device="cpu") + finally: + channel.close() + + torch.testing.assert_close(resolved, torch.ones(1)) + + +def test_tensor_channel_validates_bindings_and_rank() -> None: + channel = WorkerTensorChannel(consumer_world_size=2, timeout=1) + try: + channel.bind_producer() + with pytest.raises(ValueError, match="already has a producer"): + channel.bind_producer() + with pytest.raises(ValueError, match="expects 2 consumer ranks"): + channel.bind_consumer(1) + channel.bind_consumer(2) + with pytest.raises(ValueError, match="outside"): + channel.receive(None, rank=2, device="cpu") + finally: + channel.close() + + +def test_tensor_channel_transfers_between_independent_spawned_processes() -> None: + context = mp.get_context("spawn") + channel = WorkerTensorChannel(consumer_world_size=1, timeout=10) + metadata_queue = context.SimpleQueue() + result_queue = context.SimpleQueue() + release_event = context.Event() + producer = context.Process(target=_spawn_send, args=(channel, metadata_queue, release_event)) + consumer = context.Process(target=_spawn_receive, args=(channel, metadata_queue, result_queue, release_event)) + try: + producer.start() + consumer.start() + assert result_queue.get() == [0.0, 1.0, 2.0, 3.0] + producer.join(timeout=10) + consumer.join(timeout=10) + assert producer.exitcode == 0 + assert consumer.exitcode == 0 + finally: + release_event.set() + for process in (producer, consumer): + if process.is_alive(): + process.terminate() + process.join(timeout=2) + metadata_queue.close() + result_queue.close() + channel.close() + + +def test_parallel_workers_exchange_tensor_without_parent_materialization() -> None: + channel = WorkerTensorChannel(consumer_world_size=1, timeout=10) + producer = ParallelWorker( + _TensorProducerStage(), + tensor_output_channel=channel, + tensor_output_methods=("produce",), + ) + consumer = ParallelWorker(_TensorConsumerStage(), tensor_input_channels=(channel,)) + try: + ref, profile = producer.produce(sync=True) + assert isinstance(ref, WorkerTensorRef) + assert profile == {"seconds": 0.25} + result = consumer.consume(ref, sync=True) + torch.testing.assert_close(result, torch.arange(4, dtype=torch.float32) + 1) + finally: + producer.close() + consumer.close() + channel.close() + + +@pytest.mark.distributed +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") +def test_tensor_channel_uses_cuda_ipc_and_cross_gpu_peer_copy() -> None: + context = mp.get_context("spawn") + channel = WorkerTensorChannel(consumer_world_size=1, timeout=30) + metadata_queue = context.SimpleQueue() + result_queue = context.SimpleQueue() + release_event = context.Event() + producer = context.Process(target=_spawn_send_cuda, args=(channel, metadata_queue, release_event)) + consumer = context.Process(target=_spawn_receive_cuda, args=(channel, metadata_queue, result_queue, release_event)) + try: + producer.start() + consumer.start() + device, values = result_queue.get() + assert device == "cuda:1" + assert values == [0.0, 1.0, 2.0, 3.0] + producer.join(timeout=30) + consumer.join(timeout=30) + assert producer.exitcode == 0 + assert consumer.exitcode == 0 + finally: + release_event.set() + for process in (producer, consumer): + if process.is_alive(): + process.terminate() + process.join(timeout=2) + metadata_queue.close() + result_queue.close() + channel.close() From 3fdd12e73733ffd3366c64894f2ef10263097f2e Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Sun, 2 Aug 2026 10:09:36 +0000 Subject: [PATCH 06/11] perf(lingbot): keep session conditions on device Encode the bounded reference image once per session, distribute its latent directly to every DiT rank, and retain it for rank-local chunk slicing and mask construction. Connect non-colocated DiT and VAE stages through the generic tensor channel so large latents bypass the orchestrator.\n\nKeep scheduler timesteps and immutable RoPE frequency tables resident on worker devices, include condition memory in capacity accounting, and preserve the configured two-session capacity.\n\nVerification: 185 related unit tests and the full 957-frame four-GPU AIPerf runs passed; ruff, format check, and git diff --check passed. --- docs/en/stream_scheduler.md | 19 +++-- docs/zh/stream_scheduler.md | 25 +++++-- examples/lingbot/README.md | 5 ++ telefuser/models/lingbot_world_fast_dit.py | 13 +++- .../pipelines/lingbot_world_fast/denoising.py | 74 +++++++++++++++++-- .../pipelines/lingbot_world_fast/pipeline.py | 72 ++++++++++++++++-- .../pipelines/lingbot_world_fast/streaming.py | 10 ++- .../pipelines/lingbot_world_fast/vae_stage.py | 43 +++++++---- .../models/test_lingbot_world_fast_dit.py | 22 ++++++ .../lingbot_world_fast/test_module_loading.py | 65 +++++++++++++++- .../test_runtime_baseline.py | 7 +- .../lingbot_world_fast/test_session_cache.py | 27 +++++++ .../test_vae_stage_capacity.py | 23 +++++- 13 files changed, 356 insertions(+), 49 deletions(-) diff --git a/docs/en/stream_scheduler.md b/docs/en/stream_scheduler.md index 60282c8..c0ed809 100644 --- a/docs/en/stream_scheduler.md +++ b/docs/en/stream_scheduler.md @@ -77,21 +77,25 @@ that chunk. Other `BidirectionalService` implementations define their own cross- ## LingBot Condition Prefetch -LingBot condition encoding is independent of the corresponding control input. The session therefore keeps a fixed -lookahead of two conditions so VAE encode can overlap with earlier denoise and decode work: +LingBot encodes the bounded reference-image prefix once while initializing the session. A generic +`WorkerTensorChannel` distributes that base latent directly from the VAE encode worker to every DiT rank, where it +remains resident in the session cache. Later condition artifacts contain only `chunk_index` and `chunk_size`; each +rank slices the resident latent, repeats its tail when necessary, and constructs the first-frame mask locally. + +The session keeps a fixed lookahead of two condition metadata artifacts independently of control admission: - Session startup admits `condition[0]` and `condition[1]` when bounded ingress has capacity. - After denoise completes for chunk `i`, the session refills the window, normally with `condition[i+2]`. - `next_condition_index` and `next_control_index` maintain `0 <= next_condition_index - next_control_index <= 2`. -- If backpressure prevented prefetch, the next control and its missing encode request are admitted atomically. +- If backpressure prevented prefetch, the next control and its missing condition request are admitted atomically. Conditions and controls still join by session and sequence ID before denoise. The optimization changes scheduling, not model computation or causal cache ownership. `latency_anchor_artifact="control"` ensures condition-only prefetch does not start the control-to-output timer. -This model-specific policy sits above the generic scheduler; edge capacities continue to bound retained tensors and -session cleanup still runs through the owning actors. +This model-specific policy sits above the generic scheduler. Edge capacities bound in-flight metadata while retained +session capacity includes the resident condition latent, and cleanup still runs through the owning actors. ## Actor Ownership and Session Lifecycle @@ -126,6 +130,11 @@ LingBot uses independent `vae_encode_config` and `vae_decode_config`. Each VAE stage receives its own complete `ModelRuntimeConfig`; there is no shared VAE placement fallback. +When distributed DiT and VAE decode use different worker groups, LingBot connects their latent edge with a generic +`WorkerTensorChannel`. The denoising worker sends CUDA IPC handles directly to the decode ranks and returns only +validated tensor metadata to the scheduler. The parent process therefore retains bounded artifact ownership and +ordering without materializing the latent or allocating copies on the decode GPUs. + ## Observability and Real-Time Operation `StreamingSessionMetrics` records scheduler-observed timing and lifecycle data, including: diff --git a/docs/zh/stream_scheduler.md b/docs/zh/stream_scheduler.md index e52d59b..54e71c4 100644 --- a/docs/zh/stream_scheduler.md +++ b/docs/zh/stream_scheduler.md @@ -72,20 +72,23 @@ flowchart TB ## LingBot Condition 预取 -LingBot 的 condition encode 不依赖对应 control，因此 session 使用固定深度为 2 的 lookahead，让 VAE encode -与较早 chunk 的 denoise、decode 重叠： +LingBot 在 session 初始化时只编码一次有界的参考图前缀。通用 `WorkerTensorChannel` 将基础 latent 从 VAE +encode worker 直接分发到每个 DiT rank，并常驻在对应 session cache。后续 condition artifact 只包含 +`chunk_index` 和 `chunk_size`；每个 rank 本地切片、按需重复尾部 latent，并生成首帧 mask。 + +session 对 condition metadata 保持固定深度为 2 的 lookahead，且不依赖 control 准入： - session 启动时，在有界 ingress 有容量的前提下提交 `condition[0]` 和 `condition[1]`； - chunk `i` 的 denoise 完成后补满窗口，正常情况下提交 `condition[i+2]`； - `next_condition_index` 与 `next_control_index` 始终满足 `0 <= next_condition_index - next_control_index <= 2`； -- 若 backpressure 导致预取缺失，下一个 control 会与缺失的 encode request 原子准入。 +- 若 backpressure 导致预取缺失，下一个 control 会与缺失的 condition request 原子准入。 condition 与 control 仍按 session 和 sequence ID 在 denoise 前汇合。该优化只调整调度，不改变模型计算或 causal cache 所有权。`latency_anchor_artifact="control"` 保证单独预取 condition 不会启动 control-to-output 计时。 -这一模型专属策略位于通用 scheduler 之上；edge capacity 继续约束 tensor 保留量，session 清理仍通过 owning -actor 执行。 +这一模型专属策略位于通用 scheduler 之上；edge capacity 约束在途 metadata，常驻 condition latent 计入 +session capacity，session 清理仍通过 owning actor 执行。 ## Actor 所有权与 Session 生命周期 @@ -109,13 +112,19 @@ LingBot 的离线 chunked generation，以及通过 LiveKit 传输的双向 sess `StreamingResourceGroupSpec` 表示显式的共享并发约束。只有当 `StreamingStageSpec.resource_group` 引用 `StreamingPipelineSpec.resource_groups` 中声明的 group 时，stage 才会参与该约束。 -不要根据 `device_id` 或 `ParallelConfig.device_ids` 推断 resource group。对于 LingBot，VAE encode、DiT 和 -VAE decode 是独立 actor，即使位于同一张 GPU 也可以重叠执行。若放置超过显存容量，应移动 stage 到其他设备， -或声明明确的部署约束；不要增加隐式的全局互斥锁。 +不要根据 `device_id` 或 `ParallelConfig.device_ids` 推断 resource group。LingBot VAE encode 保持为独立 +actor。当分布式 DiT 和 VAE decode 使用完全相同的 device list 与 world size 时，pipeline 会把 decoder +co-locate 到 DiT worker group，以复用 CUDA context。若放置超过显存容量，应移动 stage 到其他设备，或声明 +明确的部署约束；不要增加隐式的全局互斥锁。 LingBot 的 `vae_encode_config` 和 `vae_decode_config` 是两个独立且完整的 `ModelRuntimeConfig`，不再提供共享的 VAE placement fallback。 +当分布式 DiT 和 VAE decode 位于不同 worker group 时，LingBot 使用通用 `WorkerTensorChannel` 连接 latent +edge。Denoising worker 直接向 decode ranks 发送 CUDA IPC handle，只把经过校验的 tensor metadata 返回给 +scheduler。主进程仍负责有界 artifact 的 ownership 和顺序，但不再 materialize latent，也不会在 decode GPU +上分配中转副本。 + ## 可观测性与实时运行 `StreamingSessionMetrics` 记录 scheduler 观测到的时序和生命周期数据，包括： diff --git a/examples/lingbot/README.md b/examples/lingbot/README.md index cb94301..30968d2 100644 --- a/examples/lingbot/README.md +++ b/examples/lingbot/README.md @@ -265,6 +265,11 @@ switching. For other counts, the examples retain the PPL-configured VAE devices Direct `LingBotWorldFastPipelineConfig` users may set `vae_encode_config`, `vae_decode_config`, and `dit_config` independently; non-matching placements continue to use independent workers. +The reference image is VAE-encoded once per session into at most 16 latent frames. For distributed DiT, the encode +worker sends that base latent once to every DiT rank through CUDA IPC/P2P; each rank retains it and builds later +four-frame condition slices and masks locally. Subsequent chunks therefore carry condition metadata rather than +repeating VAE-to-CPU-to-DiT transfers. The retained bytes are included in session-capacity accounting. + The scheduler does not infer a resource group from overlapping device IDs. VAE encode remains independent, while an exactly matching distributed DiT/VAE-decode placement uses the pipeline's explicit co-location path. See the [streaming scheduler guide](../../docs/en/stream_scheduler.md) for lifecycle guarantees. diff --git a/telefuser/models/lingbot_world_fast_dit.py b/telefuser/models/lingbot_world_fast_dit.py index a6360bb..109d497 100644 --- a/telefuser/models/lingbot_world_fast_dit.py +++ b/telefuser/models/lingbot_world_fast_dit.py @@ -505,6 +505,7 @@ def __init__( freqs = precompute_freqs_cis_3d(head_dim) self.freqs_cos = torch.cat([f.real for f in freqs], dim=-1) self.freqs_sin = torch.cat([f.imag for f in freqs], dim=-1) + self._freqs_by_device: dict[tuple[str, int | None], tuple[torch.Tensor, torch.Tensor]] = {} self.device_mesh: DeviceMesh | None = None self.usp_flag = False @@ -603,6 +604,15 @@ def _prepare_session_causal_rope( session_input_cache["causal_rope"] = (key, rope) return rope + def _rope_frequencies(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + """Keep the original-precision RoPE tables resident after first use.""" + key = (device.type, device.index) + cached = self._freqs_by_device.get(key) + if cached is None: + cached = (self.freqs_cos.to(device=device), self.freqs_sin.to(device=device)) + self._freqs_by_device[key] = cached + return cached + def patchify(self, x: torch.Tensor) -> tuple[torch.Tensor, tuple[int, int, int]]: x = x.contiguous(memory_format=torch.channels_last_3d) x = self.patch_embedding(x) @@ -672,8 +682,7 @@ def forward( control_tokens = prepared_control[0] if prepared_control is not None else None camera_modulations = prepared_control[1] if prepared_control is not None else None - freqs_cos = self.freqs_cos.to(device=x.device) - freqs_sin = self.freqs_sin.to(device=x.device) + freqs_cos, freqs_sin = self._rope_frequencies(x.device) causal_rope = self._prepare_session_causal_rope( freqs_cos, freqs_sin, diff --git a/telefuser/pipelines/lingbot_world_fast/denoising.py b/telefuser/pipelines/lingbot_world_fast/denoising.py index de3900d..a46993e 100644 --- a/telefuser/pipelines/lingbot_world_fast/denoising.py +++ b/telefuser/pipelines/lingbot_world_fast/denoising.py @@ -50,6 +50,7 @@ class _DenoisingCacheState: noise_generator: torch.Generator noise_shape: tuple[int, int, int, int, int] prompt_emb: torch.Tensor | None = None + image_condition_latent: torch.Tensor | None = None pool_slot: int | None = None projected_context_key: tuple[object, ...] | None = None prepared_control_key: tuple[object, ...] | None = None @@ -178,6 +179,7 @@ def __init__( self.empty_cache_after_call = False self._cache_registry: dict[int, _DenoisingCacheState] = {} self._cache_pool: _DenoisingCachePool | None = None + self._observed_condition_bytes = 0 self._vae_decode_stage: LingBotWorldFastVAEDecodeStage | None = None self._pending_vae_decode_latents: dict[int, deque[torch.Tensor]] = {} if model_runtime_config.parallel_config.world_size == 1 and model_runtime_config.compile_config.enabled: @@ -332,7 +334,7 @@ def estimate_session_cache_bytes(self, batch_size: int, kv_size: int, max_sequen self_kv = 2 * self.dit.num_layers * batch_size * kv_size * local_num_heads * head_dim * element_size cross_kv = 2 * self.dit.num_layers * batch_size * max_sequence_length * num_heads * head_dim * element_size cursors = self.dit.num_layers * 2 * torch.empty((), dtype=torch.int64).element_size() - return self_kv + cross_kv + cursors + return self_kv + cross_kv + cursors + getattr(self, "_observed_condition_bytes", 0) def configure_cache_pool( self, @@ -397,6 +399,7 @@ def initialize_cache( noise_generator_state: list[int], noise_shape: tuple[int, int, int, int, int], prompt_emb: torch.Tensor | None = None, + image_condition: dict[str, object] | None = None, timestep_indices: tuple[int, ...] = (0, 179, 358, 679), ) -> bool: """Atomically register session-scoped KV, scheduler, and RNG state.""" @@ -404,7 +407,7 @@ def initialize_cache( raise ValueError(f"Cache handle {cache_handle} is already registered") scheduler = FlowUniPCMultistepScheduler(num_train_timesteps=1000, shift=1, use_dynamic_shifting=False) - timesteps = _select_timesteps(scheduler, tuple(timestep_indices), sample_shift) + timesteps = _select_timesteps(scheduler, tuple(timestep_indices), sample_shift).to(self.device) generator = torch.Generator(device=self.device) generator.set_state(torch.tensor(generator_state, dtype=torch.uint8)) noise_generator = torch.Generator(device=self.device) @@ -441,6 +444,13 @@ def initialize_cache( prompt_emb=prompt_emb, pool_slot=pool_slot, ) + if image_condition is not None: + self._resolve_image_condition( + state, + image_condition, + device=self.device, + dtype=self.torch_dtype, + ) except Exception: if pool is not None and pool_slot is not None: pool.release(pool_slot) @@ -467,6 +477,54 @@ def _convert_flow_pred_to_x0( x0 = xt - sigma_t * flow_pred return x0.to(original_dtype) + def _resolve_image_condition( + self, + state: _DenoisingCacheState, + condition_chunk: torch.Tensor | dict[str, object], + *, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + """Materialize one condition chunk from a session-resident image latent.""" + if isinstance(condition_chunk, torch.Tensor): + return condition_chunk + if not isinstance(condition_chunk, dict): + raise TypeError( + f"LingBot image condition must be a tensor or mapping, got {type(condition_chunk).__name__}" + ) + + chunk_index = condition_chunk.get("chunk_index") + chunk_size = condition_chunk.get("chunk_size") + if not isinstance(chunk_index, int) or chunk_index < 0: + raise ValueError(f"LingBot image condition has invalid chunk_index {chunk_index!r}") + if not isinstance(chunk_size, int) or chunk_size < 1: + raise ValueError(f"LingBot image condition has invalid chunk_size {chunk_size!r}") + + exported = condition_chunk.get("latent_condition") + if exported is not None: + if not isinstance(exported, torch.Tensor) or exported.ndim != 4 or exported.shape[1] < 1: + raise ValueError("LingBot image condition latent must have shape (channels, frames, height, width)") + state.image_condition_latent = exported.to(device=device, dtype=dtype) + condition_bytes = state.image_condition_latent.numel() * state.image_condition_latent.element_size() + self._observed_condition_bytes = max(getattr(self, "_observed_condition_bytes", 0), condition_bytes) + latent_condition = state.image_condition_latent + if latent_condition is None: + raise RuntimeError("LingBot image condition metadata arrived before the session latent") + + start = chunk_index * chunk_size + available = latent_condition[:, start : start + chunk_size] + if available.shape[1] < chunk_size: + tail = latent_condition[:, -1:].expand(-1, chunk_size - available.shape[1], -1, -1) + available = torch.cat([available, tail], dim=1) + mask = torch.zeros( + (4, chunk_size, available.shape[2], available.shape[3]), + dtype=available.dtype, + device=available.device, + ) + if chunk_index == 0: + mask[:, 0] = 1 + return torch.cat([mask, available], dim=0).unsqueeze(0) + @staticmethod def _build_i2v_model_input_writer( latent_chunk: torch.Tensor, @@ -514,7 +572,7 @@ def denoise_chunk( step_start = torch.cuda.Event(enable_timing=True) step_end = torch.cuda.Event(enable_timing=True) step_start.record() - schedule_timestep = timesteps[timestep_idx].view(1).to(device=current_latent.device) + schedule_timestep = timesteps[timestep_idx].view(1) model_timestep = schedule_timestep.to(dtype=torch.float32) with torch.amp.autocast( current_latent.device.type, @@ -544,7 +602,7 @@ def denoise_chunk( raise RuntimeError("LingBot DMD forward unexpectedly returned no prediction") x0 = self._convert_flow_pred_to_x0(noise_pred, current_latent, schedule_timestep[0], scheduler) if timestep_idx < len(timesteps) - 1: - next_timestep = timesteps[timestep_idx + 1].view(1).to(device=x0.device) + next_timestep = timesteps[timestep_idx + 1].view(1) noise = torch.randn(x0.shape, generator=generator, device=x0.device, dtype=x0.dtype) current_latent = scheduler.add_noise(x0, noise, next_timestep) else: @@ -604,7 +662,7 @@ def _prepare_session_inputs( def denoise_and_update_cache( self, cache_handle: int, - condition_chunk: torch.Tensor, + condition_chunk: torch.Tensor | dict[str, object], prompt_emb: torch.Tensor | None, control_chunk: torch.Tensor | None, current_start: int, @@ -623,6 +681,12 @@ def denoise_and_update_cache( self._prepare_session_inputs(state, cache_handle, session_prompt_emb, control_chunk) try: latent_chunk = self._next_noise_chunk(state) + condition_chunk = self._resolve_image_condition( + state, + condition_chunk, + device=latent_chunk.device, + dtype=self.torch_dtype, + ) prepare_model_input = self._build_i2v_model_input_writer( latent_chunk, condition_chunk, diff --git a/telefuser/pipelines/lingbot_world_fast/pipeline.py b/telefuser/pipelines/lingbot_world_fast/pipeline.py index f49b324..cba5275 100644 --- a/telefuser/pipelines/lingbot_world_fast/pipeline.py +++ b/telefuser/pipelines/lingbot_world_fast/pipeline.py @@ -28,6 +28,7 @@ from telefuser.utils.logging import logger from telefuser.utils.profiler import ProfilingContext4Debug from telefuser.worker.parallel_worker import ParallelWorker +from telefuser.worker.tensor_channel import WorkerTensorChannel from .control import ( LingBotWorldFastControlBuilder, @@ -322,15 +323,30 @@ def init(self, module_manager: ModuleManager, config: LingBotWorldFastPipelineCo self.vae_decode_device = self._runtime_device(vae_decode_config) self.vae_encode_torch_dtype = vae_encode_config.torch_dtype self.vae_device = self.vae_decode_device + dit_runtime_config = config.dit_config + self._worker_tensor_channels: list[WorkerTensorChannel] = [] + condition_channel = None + if dit_runtime_config.parallel_config.world_size > 1: + condition_channel = WorkerTensorChannel( + dit_runtime_config.parallel_config.world_size, + timeout=max( + vae_encode_config.parallel_config.timeout, + dit_runtime_config.parallel_config.timeout, + ), + ) + self._worker_tensor_channels.append(condition_channel) vae_encode_stage = LingBotWorldFastVAEEncodeStage( "lingbot_world_fast_vae_encode", module_manager, vae_encode_config ) vae_decode_stage = LingBotWorldFastVAEDecodeStage( "lingbot_world_fast_vae_decode", module_manager, vae_decode_config ) - self.vae_encode_worker = ParallelWorker(vae_encode_stage) + self.vae_encode_worker = ParallelWorker( + vae_encode_stage, + tensor_output_channel=condition_channel, + tensor_output_methods=("encode_condition_chunk",) if condition_channel is not None else (), + ) - dit_runtime_config = config.dit_config dit_device = "cpu" if dit_runtime_config.parallel_config.world_size > 1 else self.device self.dit = module_manager.fetch_module("lingbot_world_fast_dit") if self.dit is None: @@ -350,12 +366,37 @@ def init(self, module_manager: ModuleManager, config: LingBotWorldFastPipelineCo ) if colocate_vae_decode: denoise_stage.attach_vae_decode_stage(vae_decode_stage) + direct_vae_handoff = dit_runtime_config.parallel_config.world_size > 1 and not colocate_vae_decode + latent_channel = None + if direct_vae_handoff: + latent_channel = WorkerTensorChannel( + vae_decode_config.parallel_config.world_size, + timeout=max( + dit_runtime_config.parallel_config.timeout, + vae_decode_config.parallel_config.timeout, + ), + ) + self._worker_tensor_channels.append(latent_channel) self.denoise_stage = ( - ParallelWorker(denoise_stage) if dit_runtime_config.parallel_config.world_size > 1 else denoise_stage + ParallelWorker( + denoise_stage, + tensor_output_channel=latent_channel, + tensor_output_methods=("denoise_and_update_cache",) if latent_channel is not None else (), + tensor_input_channels=(condition_channel,) if condition_channel is not None else (), + ) + if dit_runtime_config.parallel_config.world_size > 1 + else denoise_stage ) self.vae_decode_worker = ( - _CoLocatedVAEDecodeWorker(self.denoise_stage) if colocate_vae_decode else ParallelWorker(vae_decode_stage) + _CoLocatedVAEDecodeWorker(self.denoise_stage) + if colocate_vae_decode + else ParallelWorker( + vae_decode_stage, + tensor_input_channels=(latent_channel,) if latent_channel is not None else (), + ) ) + self.uses_direct_condition_handoff = condition_channel is not None + self.uses_direct_vae_handoff = direct_vae_handoff @staticmethod def _validate_vae_stage_runtime_config( @@ -547,8 +588,8 @@ def _prepare_image_tensor(self, image: Image.Image, height: int, width: int) -> encode_dtype = getattr(self, "vae_encode_torch_dtype", self.config.vae_encode_config.torch_dtype) return tensor.to("cpu", dtype=encode_dtype) - def _initialize_vae_session(self, session: LingBotWorldFastGenerationSession) -> None: - """Register session-owned VAE caches in the dedicated worker.""" + def _initialize_vae_session(self, session: LingBotWorldFastGenerationSession) -> dict[str, object]: + """Register VAE caches and export the session image latent once.""" if session.cache_handle is None or session.condition_image is None: raise RuntimeError("VAE session initialization requires an image and cache handle") encoder_initialized = self.vae_encode_worker.initialize_cache( @@ -559,6 +600,20 @@ def _initialize_vae_session(self, session: LingBotWorldFastGenerationSession) -> decoder_initialized = self.vae_decode_worker.initialize_cache(cache_handle=session.cache_handle, sync=True) if not decoder_initialized: raise RuntimeError("LingBot retained-session cache capacity is exhausted at VAE decode") + condition = self.vae_encode_worker.encode_condition_chunk( + cache_handle=session.cache_handle, + chunk_index=0, + chunk_count=session.chunk_count, + chunk_size=session.chunk_size, + height=session.height, + width=session.width, + output_dtype=self.torch_dtype, + sync=True, + ) + if not isinstance(condition, dict): + raise TypeError("LingBot VAE condition initialization must return a mapping") + session.condition_image = None + return condition def _release_vae_session_cache(self, session: LingBotWorldFastGenerationSession) -> bool: if not hasattr(self, "vae_encode_worker") or not hasattr(self, "vae_decode_worker"): @@ -613,6 +668,8 @@ def close(self) -> None: for vae_worker in (getattr(self, "vae_encode_worker", None), getattr(self, "vae_decode_worker", None)): if isinstance(vae_worker, ParallelWorker): vae_worker.close() + for channel in getattr(self, "_worker_tensor_channels", ()): + channel.close() def _get_streaming_runtime(self) -> LingBotWorldFastStreamingRuntime: """Return the one actor graph owned by this pipeline instance.""" @@ -723,7 +780,7 @@ def _create_initialized_session( cache_handle=cache_handle, ) try: - self._initialize_vae_session(session) + image_condition = self._initialize_vae_session(session) initialize_cache_kwargs = dict( cache_handle=cache_handle, batch_size=1, @@ -734,6 +791,7 @@ def _create_initialized_session( noise_generator_state=noise_generator.get_state().tolist(), noise_shape=(1, 16, session_config.chunk_size, lat_h, lat_w), prompt_emb=prompt_emb, + image_condition=image_condition, timestep_indices=getattr(self.config, "timestep_indices", (0, 179, 358, 679)), ) if isinstance(self.denoise_stage, ParallelWorker): diff --git a/telefuser/pipelines/lingbot_world_fast/streaming.py b/telefuser/pipelines/lingbot_world_fast/streaming.py index 07cc005..7c1d828 100644 --- a/telefuser/pipelines/lingbot_world_fast/streaming.py +++ b/telefuser/pipelines/lingbot_world_fast/streaming.py @@ -450,20 +450,21 @@ def _encode_inputs(self, invocation: StreamingStageInvocation) -> tuple[tuple[ob "chunk_size": runtime.chunk_size, "height": runtime.height, "width": runtime.width, + "output_dtype": self.pipeline.torch_dtype, } - def _encode_outputs(self, value: torch.Tensor, invocation: StreamingStageInvocation) -> dict[str, object]: + def _encode_outputs(self, value: dict[str, object], invocation: StreamingStageInvocation) -> dict[str, object]: entry = self._entry_for_invocation(invocation) index = invocation.key.sequence_id self.pipeline._notify_progress(entry.progress_callback, "condition_chunk_encoded", index=index) if index == 0: entry.runtime.condition_image = None - return {"condition": value.to(device=self.pipeline.device, dtype=self.pipeline.torch_dtype)} + return {"condition": value} def _denoise_kwargs(self, invocation: StreamingStageInvocation) -> dict[str, object]: runtime = self._entry_for_invocation(invocation).runtime index = invocation.key.sequence_id - return { + kwargs = { "cache_handle": runtime.cache_handle, "condition_chunk": invocation.inputs["condition"], "prompt_emb": None, @@ -476,6 +477,9 @@ def _denoise_kwargs(self, invocation: StreamingStageInvocation) -> dict[str, obj ), "_benchmark_profile": runtime.config.benchmark_metrics, } + if getattr(self.pipeline, "uses_direct_vae_handoff", False): + kwargs["_tensor_transport"] = runtime.world_kv_binding is None + return kwargs @torch.inference_mode() def _denoise(self, invocation: StreamingStageInvocation) -> dict[str, object]: diff --git a/telefuser/pipelines/lingbot_world_fast/vae_stage.py b/telefuser/pipelines/lingbot_world_fast/vae_stage.py index e15e3d5..270eb58 100644 --- a/telefuser/pipelines/lingbot_world_fast/vae_stage.py +++ b/telefuser/pipelines/lingbot_world_fast/vae_stage.py @@ -106,6 +106,7 @@ class _VAEEncodeCacheState: condition_image: torch.Tensor | None encoder_state: WanVideoVAEStreamingEncodeState = field(default_factory=WanVideoVAEStreamingEncodeState) latent_condition: torch.Tensor | None = None + latent_condition_exported: bool = False pool_slot: int | None = None @@ -123,6 +124,7 @@ def __init__(self, name: str, module_manager: ModuleManager, model_runtime_confi self.empty_cache_after_call = False self._cache_registry: dict[int, _VAEEncodeCacheState] = {} self._observed_session_cache_bytes = 0 + self._observed_condition_bytes = 0 self._cache_layout: dict[int, tuple[torch.dtype, int]] = {} self._cache_pool: _VAECachePool | None = None @@ -138,7 +140,7 @@ def _observe_cache(self, cache: list[object]) -> None: def estimate_session_cache_bytes(self) -> int: """Return VAE encoder bytes for one fixed slot including shape headroom.""" - return sum( + causal_cache_bytes = sum( ( (numel * _CACHE_POOL_HEADROOM_NUMERATOR + _CACHE_POOL_HEADROOM_DENOMINATOR - 1) // _CACHE_POOL_HEADROOM_DENOMINATOR @@ -146,6 +148,7 @@ def estimate_session_cache_bytes(self) -> int: * torch.empty((), dtype=dtype).element_size() for dtype, numel in self._cache_layout.values() ) + return causal_cache_bytes + getattr(self, "_observed_condition_bytes", 0) def configure_cache_pool(self, capacity: int) -> VAECachePoolProfile: """Allocate all persistent encoder-cache slots before accepting sessions.""" @@ -197,9 +200,16 @@ def initialize_cache(self, cache_handle: int, condition_image: torch.Tensor) -> @with_model_offload(["vae"]) def encode_condition_chunk( - self, cache_handle: int, chunk_index: int, chunk_count: int, chunk_size: int, height: int, width: int - ) -> torch.Tensor: - """Encode the bounded zero-frame prefix once and return one condition chunk.""" + self, + cache_handle: int, + chunk_index: int, + chunk_count: int, + chunk_size: int, + height: int, + width: int, + output_dtype: torch.dtype, + ) -> dict[str, object]: + """Encode the bounded image prefix once and export it once per session.""" state = self._cache_registry[cache_handle] if state.latent_condition is None: if state.condition_image is None: @@ -223,19 +233,22 @@ def encode_condition_chunk( raise RuntimeError( f"VAE condition prefix has {latent.shape[1]} latent frames, expected {encoded_latent_frames}" ) - state.latent_condition = latent.cpu() + state.latent_condition = latent.to(dtype=output_dtype) + self._observed_condition_bytes = max( + getattr(self, "_observed_condition_bytes", 0), + state.latent_condition.numel() * state.latent_condition.element_size(), + ) state.condition_image = None - latent_condition = state.latent_condition - start = chunk_index * chunk_size - available = latent_condition[:, start : start + chunk_size] - if available.shape[1] < chunk_size: - tail = latent_condition[:, -1:].expand(-1, chunk_size - available.shape[1], -1, -1) - available = torch.cat([available, tail], dim=1) - mask = torch.zeros((4, chunk_size, available.shape[2], available.shape[3]), dtype=available.dtype) - if chunk_index == 0: - mask[:, 0] = 1 - return torch.cat([mask, available], dim=0).unsqueeze(0) + exported = None + if not state.latent_condition_exported: + exported = state.latent_condition + state.latent_condition_exported = True + return { + "chunk_index": chunk_index, + "chunk_size": chunk_size, + "latent_condition": exported, + } def release_cache(self, cache_handle: int) -> bool: """Release encoder state for one session.""" diff --git a/tests/unit/models/test_lingbot_world_fast_dit.py b/tests/unit/models/test_lingbot_world_fast_dit.py index f7cc5a1..2fb6e07 100644 --- a/tests/unit/models/test_lingbot_world_fast_dit.py +++ b/tests/unit/models/test_lingbot_world_fast_dit.py @@ -161,6 +161,28 @@ def test_set_attention_config_updates_all_blocks() -> None: assert block.cross_attn.attention_config is attention_config +def test_rope_frequencies_keep_original_precision_and_are_cached_per_device() -> None: + model = LingBotWorldFastDiT( + in_dim=4, + dim=32, + ffn_dim=64, + freq_dim=8, + text_dim=16, + out_dim=4, + num_heads=4, + num_layers=1, + ).to(dtype=torch.bfloat16) + + original_dtype = model.freqs_cos.dtype + first = model._rope_frequencies(torch.device("cpu")) + second = model._rope_frequencies(torch.device("cpu")) + + assert first[0].dtype == original_dtype + assert first[1].dtype == original_dtype + assert second[0] is first[0] + assert second[1] is first[1] + + def test_scalar_timestep_modulation_stays_broadcastable() -> None: model = LingBotWorldFastDiT( in_dim=4, diff --git a/tests/unit/pipelines/lingbot_world_fast/test_module_loading.py b/tests/unit/pipelines/lingbot_world_fast/test_module_loading.py index fa2e410..323db05 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_module_loading.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_module_loading.py @@ -2,7 +2,7 @@ import torch -from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.config import ModelRuntimeConfig, ParallelConfig from telefuser.pipelines.lingbot_world_fast.pipeline import LingBotWorldFastPipeline, LingBotWorldFastPipelineConfig @@ -47,3 +47,66 @@ def fetch_module(name: str, require_model_path: bool = False): assert module_manager.fetch_module.call_args_list[0].kwargs == {"require_model_path": True} assert module_manager.fetch_module.call_args_list[1].args == ("lingbot_world_fast_dit",) dit.set_causal_attention_window.assert_called_once_with(-1, 0) + + +def test_non_colocated_dit_and_vae_workers_use_direct_tensor_channel(tmp_path) -> None: + text_encoder = MagicMock() + dit = MagicMock(control_type="cam") + text_encoder.to.return_value = text_encoder + text_encoder.eval.return_value = text_encoder + dit.to.return_value = dit + dit.eval.return_value = dit + dit.requires_grad_.return_value = dit + module_manager = MagicMock() + module_manager.fetch_module.side_effect = lambda name, require_model_path=False: ( + (text_encoder, str(tmp_path / "models_t5_umt5-xxl-enc-bf16.pth")) + if name == "wan_video_text_encoder" and require_model_path + else dit + if name == "lingbot_world_fast_dit" + else None + ) + condition_channel = MagicMock() + latent_channel = MagicMock() + + with ( + patch("telefuser.pipelines.lingbot_world_fast.pipeline.ParallelWorker") as worker_cls, + patch( + "telefuser.pipelines.lingbot_world_fast.pipeline.WorkerTensorChannel", + side_effect=(condition_channel, latent_channel), + ) as channel_cls, + patch("telefuser.pipelines.lingbot_world_fast.pipeline.LingBotWorldFastVAEEncodeStage"), + patch("telefuser.pipelines.lingbot_world_fast.pipeline.LingBotWorldFastVAEDecodeStage"), + patch("telefuser.pipelines.lingbot_world_fast.pipeline.LingBotWorldFastDenoisingStage"), + patch("telefuser.pipelines.lingbot_world_fast.pipeline.HuggingfaceTokenizer"), + ): + pipeline = LingBotWorldFastPipeline(device="cpu", torch_dtype=torch.float32) + pipeline.init( + module_manager, + LingBotWorldFastPipelineConfig( + dit_config=ModelRuntimeConfig( + torch_dtype=torch.float32, + parallel_config=ParallelConfig(device_ids=[0, 1], sp_ulysses_degree=2), + ), + vae_decode_config=ModelRuntimeConfig( + torch_dtype=torch.float32, + parallel_config=ParallelConfig(device_ids=[2]), + ), + ), + ) + + assert channel_cls.call_args_list[0].args == (2,) + assert channel_cls.call_args_list[0].kwargs == {"timeout": 600} + assert channel_cls.call_args_list[1].args == (1,) + assert channel_cls.call_args_list[1].kwargs == {"timeout": 600} + assert pipeline.uses_direct_condition_handoff is True + assert pipeline.uses_direct_vae_handoff is True + assert worker_cls.call_args_list[0].kwargs == { + "tensor_output_channel": condition_channel, + "tensor_output_methods": ("encode_condition_chunk",), + } + assert worker_cls.call_args_list[1].kwargs == { + "tensor_output_channel": latent_channel, + "tensor_output_methods": ("denoise_and_update_cache",), + "tensor_input_channels": (condition_channel,), + } + assert worker_cls.call_args_list[2].kwargs == {"tensor_input_channels": (latent_channel,)} diff --git a/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py b/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py index bda0b4e..d981017 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py @@ -37,6 +37,11 @@ def _build_runtime_pipeline() -> LingBotWorldFastPipeline: ) pipeline.denoise_stage = MagicMock() pipeline.vae_encode_worker = MagicMock() + pipeline.vae_encode_worker.encode_condition_chunk.return_value = { + "chunk_index": 0, + "chunk_size": 3, + "latent_condition": torch.zeros(16, 3, 2, 2), + } pipeline.vae_decode_worker = MagicMock() pipeline._next_cache_handle = 0 pipeline.encode_prompt = MagicMock(return_value=torch.zeros(1, 4, 8)) @@ -168,7 +173,7 @@ def test_aligned_81_frame_runtime_has_seven_complete_latent_chunks() -> None: assert runtime.latent_f == 21 assert runtime.chunk_count == 7 assert not hasattr(runtime, "noise_generator") - assert runtime.condition_image is not None + assert runtime.condition_image is None assert not hasattr(runtime, "noise_chunks") assert not hasattr(runtime, "condition_chunks") assert runtime.cache_handle == 0 diff --git a/tests/unit/pipelines/lingbot_world_fast/test_session_cache.py b/tests/unit/pipelines/lingbot_world_fast/test_session_cache.py index 38e0f99..2032b2f 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_session_cache.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_session_cache.py @@ -47,6 +47,33 @@ def test_worker_cache_retains_session_prompt_embedding() -> None: assert _initialize_cache(stage, 11, prompt_emb=prompt_emb) is True assert stage._cache_registry[11].prompt_emb is prompt_emb + assert stage._cache_registry[11].timesteps.device == stage.device + + +def test_worker_cache_retains_image_condition_and_materializes_chunks_locally() -> None: + stage = _cache_stage() + _initialize_cache(stage, 11) + state = stage._cache_registry[11] + latent = torch.arange(16 * 5, dtype=torch.float32).view(16, 5, 1, 1) + + first = stage._resolve_image_condition( + state, + {"chunk_index": 0, "chunk_size": 2, "latent_condition": latent}, + device=stage.device, + dtype=torch.float32, + ) + tail = stage._resolve_image_condition( + state, + {"chunk_index": 3, "chunk_size": 2, "latent_condition": None}, + device=stage.device, + dtype=torch.float32, + ) + + assert state.image_condition_latent is latent + assert first.shape == (1, 20, 2, 1, 1) + assert torch.equal(first[0, :4, 0], torch.ones(4, 1, 1)) + assert torch.count_nonzero(first[0, :4, 1]) == 0 + assert torch.equal(tail[0, 4:], latent[:, -1:].expand(-1, 2, -1, -1)) def test_worker_cache_registry_isolates_handles_and_releases_idempotently() -> None: diff --git a/tests/unit/pipelines/lingbot_world_fast/test_vae_stage_capacity.py b/tests/unit/pipelines/lingbot_world_fast/test_vae_stage_capacity.py index 9988fa3..b34b0dd 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_vae_stage_capacity.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_vae_stage_capacity.py @@ -7,6 +7,7 @@ import torch from telefuser.pipelines.lingbot_world_fast import vae_stage +from telefuser.pipelines.lingbot_world_fast.denoising import LingBotWorldFastDenoisingStage from telefuser.pipelines.lingbot_world_fast.vae_stage import ( _VAECachePool, _VAEDecodeCacheState, @@ -134,10 +135,28 @@ def test_vae_encode_stage_encodes_bounded_prefix_once_and_repeats_tail() -> None assert stage.initialize_cache(1, torch.ones(3, 2, 2)) is True encode = vae_stage.LingBotWorldFastVAEEncodeStage.encode_condition_chunk.__wrapped__ - first = encode(stage, 1, 0, 5, 4, 2, 2) - tail = encode(stage, 1, 4, 5, 4, 2, 2) + first_packet = encode(stage, 1, 0, 5, 4, 2, 2, torch.float32) + tail_packet = encode(stage, 1, 4, 5, 4, 2, 2, torch.float32) + resolver = LingBotWorldFastDenoisingStage.__new__(LingBotWorldFastDenoisingStage) + resolver._observed_condition_bytes = 0 + denoise_state = SimpleNamespace(image_condition_latent=None) + first = resolver._resolve_image_condition( + denoise_state, + first_packet, + device=torch.device("cpu"), + dtype=torch.float32, + ) + tail = resolver._resolve_image_condition( + denoise_state, + tail_packet, + device=torch.device("cpu"), + dtype=torch.float32, + ) assert stage.vae.frame_counts == [61] + assert first_packet["latent_condition"].shape == (16, 16, 2, 2) + assert tail_packet["latent_condition"] is None + assert denoise_state.image_condition_latent is first_packet["latent_condition"] assert first.shape == (1, 20, 4, 2, 2) assert torch.equal(first[0, :4, 0], torch.ones(4, 2, 2)) assert torch.count_nonzero(first[0, :4, 1:]) == 0 From cec5df1155a83ba7d167c06689cbf2621e526915 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Sun, 2 Aug 2026 10:10:21 +0000 Subject: [PATCH 07/11] docs(benchmark): record aligned LingBot results Document the complete 957-frame TeleFuser runs in the cu128 and cu130 environments and compare target-side compute against the retained same-environment SGLang artifact. Keep client stream throughput separate because LiveKit uses realtime pacing while the WebSocket reference uses burst output.\n\nRecord synchronized compute time, first-frame decomposition, transport boundaries, warmup semantics, and artifact identifiers in both English and Chinese benchmark guides.\n\nVerification: both documented TeleFuser sessions completed 60 chunks successfully; git diff --check passed. --- docs/en/benchmark_aiperf.md | 30 ++++++++++++++++++++++++++++++ docs/zh/benchmark_aiperf.md | 26 ++++++++++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/docs/en/benchmark_aiperf.md b/docs/en/benchmark_aiperf.md index f474395..f0cf212 100644 --- a/docs/en/benchmark_aiperf.md +++ b/docs/en/benchmark_aiperf.md @@ -102,6 +102,36 @@ Target facts follow these rules: Client delivery, target pipeline residence, target phase time, and resource utilization remain separate dimensions. Fields without equivalent semantics remain private or unavailable instead of being forced into a common metric. +## Validated one-minute LingBot-World v2 replay + +Commit `663c385b179012c5c3de613212d10e8e6eac5f5d` was validated on 2026-08-02 with the +`stream_lingbot_world_v2_1min.json` workload, AIPerf 0.11.0 at commit +`e977ffbb1648510acec431b2a3fbd1a0f7bb8a35`, and four H100 80 GB GPUs. The current H100 example used BF16 DiT, +FP32 VAE, FlashAttention-4, disabled `torch.compile`, disabled FSDP, `chunk_size=4`, and 16 FPS. The 60-second +request was truncated to 60 complete latent chunks: 957 generated frames representing 59.75 seconds of media. +LingBot-World v2 used `local_attn_size=18` and `sink_size=6`; the session reported a fixed 28,080-token KV capacity +for its 240 latent frames. + +| Runtime / target | Compute FPS | Mean / p99 chunk | Stream FPS | Client frames | Artifact | +|---|---:|---:|---:|---:|---| +| TeleFuser `.venv`, torch cu128 | 16.191 | 0.988 / 1.099 s | 12.697 | 756 | `20260802_084922_d7ae0931` | +| TeleFuser `.venv-sglang`, torch cu130 | 15.897 | 1.006 / 1.208 s | 14.089 | 871 | `20260802_090301_af6c433c` | +| SGLang `.venv-sglang`, torch cu130 | 16.617 | 0.963 / 0.974 s | 16.772 | 957 | `20260801_104829_2320fd7f` | + +Every row completed 60 target chunks and generated 957 frames. AIPerf excluded only target chunk 0, leaving 944 +frames across 59 chunks. The aligned TeleFuser run used 59.381647 seconds of synchronized compute time and was 4.33% +below the SGLang compute rate. The cu130 TeleFuser result was 1.81% below its cu128 run, so the environment change is +reported separately and is not counted as an optimization gain. The aligned TeleFuser report is +`artifacts/telefuser_aiperf/stream_lingbot_v2_1min/20260802_090301_af6c433c/stream_report.html`. + +`stream_fps` is not used for the compute comparison. TeleFuser published LiveKit video with real-time 16 FPS pacing; +its aligned run averaged 18.99 ms from decoded-ready to publish start, 941.66 ms in paced publication, and 2.10 ms +from publish completion to client metadata. SGLang used unpaced burst WebSocket output. Those delivery semantics are +not equivalent even though both include network and client decoding. TeleFuser's 9.740-second first-frame latency +comprised 0.630 seconds to create the session, another 1.979 seconds to connect, 3.206 seconds from connection to +admission, and 3.925 seconds from admission to the first client frame; runtime creation occupied 1.564 seconds of the +last interval. + ## Reproducibility Every result should retain the TeleFuser commit and AIPerf package version, model revision, accelerator model/count, diff --git a/docs/zh/benchmark_aiperf.md b/docs/zh/benchmark_aiperf.md index 0163b37..26e1c65 100644 --- a/docs/zh/benchmark_aiperf.md +++ b/docs/zh/benchmark_aiperf.md @@ -97,6 +97,32 @@ Target 原始事实遵守以下规则： 客户端交付、target pipeline residence、target phase time 和资源利用率保持为不同维度。无法等价的字段保留为 private 或 unavailable，不强行映射为同一指标。 +## LingBot-World v2 一分钟回放实测 + +2026-08-02 使用 4 张 H100 80 GB 验证了 TeleFuser commit +`663c385b179012c5c3de613212d10e8e6eac5f5d` 和 `stream_lingbot_world_v2_1min.json` workload；AIPerf 为 +0.11.0、commit `e977ffbb1648510acec431b2a3fbd1a0f7bb8a35`。当前 H100 example 使用 BF16 DiT、FP32 VAE、 +FlashAttention-4，关闭 `torch.compile` 和 FSDP，`chunk_size=4`，输出 16 FPS。60 秒请求按完整 latent +chunk 截断为 60 个 chunk、957 帧，对应 59.75 秒媒体时长。LingBot-World v2 使用 +`local_attn_size=18`、`sink_size=6`，本次 240 latent frame session 报告的固定 KV 容量为 28,080 token。 + +| 运行环境 / target | Compute FPS | Chunk mean / p99 | Stream FPS | 客户端帧数 | Artifact | +|---|---:|---:|---:|---:|---| +| TeleFuser `.venv`，torch cu128 | 16.191 | 0.988 / 1.099 s | 12.697 | 756 | `20260802_084922_d7ae0931` | +| TeleFuser `.venv-sglang`，torch cu130 | 15.897 | 1.006 / 1.208 s | 14.089 | 871 | `20260802_090301_af6c433c` | +| SGLang `.venv-sglang`，torch cu130 | 16.617 | 0.963 / 0.974 s | 16.772 | 957 | `20260801_104829_2320fd7f` | + +三次运行均完成 60 个 target chunk、生成 957 帧。AIPerf 只排除 target chunk 0，稳态统计包含 59 个 chunk、 +944 帧。对齐环境后的 TeleFuser 同步计算时间为 59.381647 秒，compute FPS 比 SGLang 低 4.33%。TeleFuser +cu130 结果比 cu128 低 1.81%，因此环境变化单独报告，不计作代码优化收益。对齐环境 TeleFuser 报告位于 +`artifacts/telefuser_aiperf/stream_lingbot_v2_1min/20260802_090301_af6c433c/stream_report.html`。 + +Compute 对比不使用 `stream_fps`。TeleFuser 的 LiveKit 视频按实时 16 FPS pacing 发布；对齐环境运行中， +decoded-ready 到 publish start 平均 18.99 ms，paced publish 平均 941.66 ms，publish 完成到客户端 metadata +平均 2.10 ms。SGLang 使用无 pacing 的 WebSocket burst 输出，两者交付语义不等价，尽管两边都包含网络 +传输和客户端解码。TeleFuser 首帧为 9.740 秒：session 创建 0.630 秒，其后连接 1.979 秒，连接到准入 +3.206 秒，准入到客户端首帧 3.925 秒；最后一段中的 runtime creation 为 1.564 秒。 + ## 复现要求 每个结果都应保留 TeleFuser commit、AIPerf 包版本、模型 revision、加速器型号/数量、driver、CUDA、 From 7b1a8f2baa990c1e23cbc6c3b9e1cccfd38726b3 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Mon, 3 Aug 2026 01:40:38 +0000 Subject: [PATCH 08/11] fix(regression): enforce SDPA and repair tensor handoff Force example regression runs onto TORCH_SDPA and disable Diffusers' optional xformers import path so the test backend is deterministic. Route Wan2.2 VAE and denoising tensors through WorkerTensorChannel, add lifecycle coverage, and document the regression backend. Verified with focused unit and worker tests plus the four previously failing GPU regression cases. --- docs/en/testing.md | 3 + docs/zh/testing.md | 3 + examples/README.md | 2 + examples/run_examples.py | 21 ++++- telefuser/pipelines/wan_video/wan22_video.py | 70 +++++++++++++-- .../pipelines/wan_video/test_wan22_video.py | 87 +++++++++++++++++++ tests/unit/test_run_examples.py | 11 +++ 7 files changed, 188 insertions(+), 9 deletions(-) create mode 100644 tests/unit/pipelines/wan_video/test_wan22_video.py diff --git a/docs/en/testing.md b/docs/en/testing.md index c565f69..307da60 100644 --- a/docs/en/testing.md +++ b/docs/en/testing.md @@ -221,6 +221,9 @@ class TestRMSNorm: TeleFuser provides a batch regression testing framework for running example pipelines, comparing outputs against baselines, and generating reports. +All regression examples run with `TORCH_SDPA`. The runner disables Diffusers' optional xformers path before importing +an example so an installed FlashAttention or xformers package cannot change the regression backend. + ### Quick Start ```bash diff --git a/docs/zh/testing.md b/docs/zh/testing.md index b847eda..10f68e4 100644 --- a/docs/zh/testing.md +++ b/docs/zh/testing.md @@ -221,6 +221,9 @@ class TestRMSNorm: TeleFuser 提供批量回归测试框架，用于运行示例 pipeline、对比 baseline 输出、生成测试报告。 +所有回归示例统一使用 `TORCH_SDPA`。runner 会在导入示例前禁用 Diffusers 的可选 xformers 路径，避免已安装的 +FlashAttention 或 xformers 包改变回归后端。 + ### 快速开始 ```bash diff --git a/examples/README.md b/examples/README.md index c8607d8..370efac 100644 --- a/examples/README.md +++ b/examples/README.md @@ -2,6 +2,8 @@ Runs configured pipelines in isolated subprocesses, compares outputs against baselines (PSNR/SSIM for video, pixel diff for image), and prints a results table. +Regression runs force every example to use `TORCH_SDPA` and disable Diffusers' +optional xformers path so results do not depend on installed attention kernels. ## Quick Start diff --git a/examples/run_examples.py b/examples/run_examples.py index 1ed6a64..051b758 100644 --- a/examples/run_examples.py +++ b/examples/run_examples.py @@ -41,6 +41,7 @@ _CONFIG_PATH = Path(__file__).resolve().parent / "example_config.yaml" _RESULT_MARKER = "###RESULT###" +_REGRESSION_ATTN_IMPL = "TORCH_SDPA" def _pipeline_slug(pipeline_key: str) -> str: @@ -780,6 +781,19 @@ def _patch_ppl_config(module: ModuleType, overrides: dict) -> None: config[key] = value +def _prepare_sdpa_regression(overrides: dict) -> dict: + """Force regression examples onto SDPA without importing optional xformers kernels.""" + try: + from diffusers.utils import import_utils + + # Diffusers checks package metadata, which can mark an incompatible + # xformers installation as available and then import it eagerly. + import_utils._xformers_available = False + except ImportError: + pass + return {**overrides, "attn_impl": _REGRESSION_ATTN_IMPL} + + def _extract_click_default(module: ModuleType, param_name: str) -> str | None: """Extract a click option's default value from the module's main() command.""" import click @@ -1103,6 +1117,7 @@ def _run_single(pipeline_key: str, config_path: str | None, output_dir: str | No timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + ppl_config_overrides = _prepare_sdpa_regression(ppl_cfg.ppl_config_overrides) runner_config = { "gpu_count": ppl_cfg.gpu_count, "seed": ppl_cfg.seed, @@ -1118,11 +1133,11 @@ def _run_single(pipeline_key: str, config_path: str | None, output_dir: str | No "first_image_path": ppl_cfg.first_image_path, "last_image_path": ppl_cfg.last_image_path, "input_video_path": ppl_cfg.input_video_path, - "ppl_config_overrides": ppl_cfg.ppl_config_overrides, + "ppl_config_overrides": ppl_config_overrides, "script": ppl_cfg.script, # Pass script for filename generation } # Merge ppl_config_overrides into runner_config for height/width access - runner_config.update(ppl_cfg.ppl_config_overrides) + runner_config.update(ppl_config_overrides) pipeline = None gpu_mem_peak = 0.0 @@ -1131,7 +1146,7 @@ def _run_single(pipeline_key: str, config_path: str | None, output_dir: str | No # Phase 1: Model Loading try: module = _import_example_module(script_path) - _patch_ppl_config(module, ppl_cfg.ppl_config_overrides) + _patch_ppl_config(module, ppl_config_overrides) pipeline = _call_get_pipeline(module, runner_config) except Exception as e: tb = traceback.format_exc() diff --git a/telefuser/pipelines/wan_video/wan22_video.py b/telefuser/pipelines/wan_video/wan22_video.py index 1b9d511..9b13f86 100644 --- a/telefuser/pipelines/wan_video/wan22_video.py +++ b/telefuser/pipelines/wan_video/wan22_video.py @@ -15,6 +15,7 @@ from telefuser.utils.logging import logger from telefuser.worker.parallel_worker import ParallelWorker from telefuser.worker.ray_worker import create_ray_worker +from telefuser.worker.tensor_channel import WorkerTensorChannel from ..common.rift_vfi import RiftVFIStage from .moe_dit_denoising import MoeDitDenoisingStage @@ -87,13 +88,44 @@ def init(self, module_manager: ModuleManager, config: Wan22VideoPipelineConfig): self.text_encoding_stage = TextEncodingStage("text_encoding", module_manager, config.text_encoding_config) if config.enable_vfi: self.vfi_stage = RiftVFIStage("vfi", module_manager, config.vfi_config) + self._worker_tensor_channels: list[WorkerTensorChannel] = [] + vae_to_denoise_channel = None + denoise_to_vae_channel = None + direct_tensor_handoff = ( + config.enable_vae_parallel and config.enable_denoising_parallel and not config.enable_vae_ray + ) + if direct_tensor_handoff: + timeout = max( + config.vae_config.parallel_config.timeout, + config.dit_high_config.parallel_config.timeout, + ) + vae_to_denoise_channel = WorkerTensorChannel( + config.dit_high_config.parallel_config.world_size, + timeout=timeout, + ) + denoise_to_vae_channel = WorkerTensorChannel( + config.vae_config.parallel_config.world_size, + timeout=timeout, + ) + self._worker_tensor_channels.extend((vae_to_denoise_channel, denoise_to_vae_channel)) if config.enable_vae_ray: logger.info("enable ray actor for vae") self.vae_stage = create_ray_worker(self.vae_stage, self.config.enable_vae_parallel) elif config.enable_vae_parallel: - self.vae_stage = ParallelWorker(self.vae_stage) + self.vae_stage = ParallelWorker( + self.vae_stage, + tensor_output_channel=vae_to_denoise_channel, + tensor_output_methods=("process",) if vae_to_denoise_channel is not None else (), + tensor_input_channels=(denoise_to_vae_channel,) if denoise_to_vae_channel is not None else (), + ) if config.enable_denoising_parallel: - self.denoise_stage = ParallelWorker(self.denoise_stage) + self.denoise_stage = ParallelWorker( + self.denoise_stage, + tensor_output_channel=denoise_to_vae_channel, + tensor_output_methods=("process",) if denoise_to_vae_channel is not None else (), + tensor_input_channels=(vae_to_denoise_channel,) if vae_to_denoise_channel is not None else (), + ) + self._uses_direct_tensor_handoff = direct_tensor_handoff # Auto-enable metrics if configured if config.enable_metrics: @@ -167,6 +199,11 @@ def __call__( num_frames, **tiler_kwargs, is_ray=self.config.enable_vae_ray, + **( + {"_tensor_transport": self._uses_direct_tensor_handoff} + if isinstance(self.vae_stage, ParallelWorker) + else {} + ), ) ref_latent = ref_latent_handler() prompt_emb_list = prompt_emb_list_handler() @@ -186,6 +223,11 @@ def __call__( sigma_shift, boundary, latent_data=latent_data, + **( + {"_tensor_transport": self._uses_direct_tensor_handoff} + if isinstance(self.denoise_stage, ParallelWorker) + else {} + ), ) denoise_result = latents_handler() @@ -195,7 +237,13 @@ def __call__( latents, latent_payload = denoise_result else: latents = denoise_result - frames_handler = auto_async_call(self.vae_stage.process, "decode_video", latents, **tiler_kwargs) + frames_handler = auto_async_call( + self.vae_stage.process, + "decode_video", + latents, + **tiler_kwargs, + **({"_tensor_transport": False} if isinstance(self.vae_stage, ParallelWorker) else {}), + ) frames = frames_handler() frames = self.tensor2video(frames[0]) @@ -211,10 +259,20 @@ def __call__( return frames, latent_payload return frames + def close(self) -> None: + """Close multiprocessing workers before releasing tensor channels.""" + for stage_name in ("denoise_stage", "vae_stage"): + stage = getattr(self, stage_name, None) + if isinstance(stage, ParallelWorker): + stage.close() + for channel in getattr(self, "_worker_tensor_channels", ()): + channel.close() + def __del__(self): - del self.vae_stage - del self.denoise_stage - del self.text_encoding_stage + try: + self.close() + except Exception: + pass @classmethod def from_pretrained( diff --git a/tests/unit/pipelines/wan_video/test_wan22_video.py b/tests/unit/pipelines/wan_video/test_wan22_video.py new file mode 100644 index 0000000..68fd302 --- /dev/null +++ b/tests/unit/pipelines/wan_video/test_wan22_video.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +from unittest.mock import Mock + +import torch + +from telefuser.core.config import ModelRuntimeConfig, ParallelConfig +from telefuser.pipelines.wan_video import wan22_video + + +class _FakeStage: + def __init__(self, _name: str, _module_manager: object, runtime_config: ModelRuntimeConfig, *_args: object) -> None: + self.model_runtime_config = runtime_config + + +class _FakeWorker: + instances: list["_FakeWorker"] = [] + + def __init__(self, stage: _FakeStage, **kwargs: object) -> None: + self.stage = stage + self.kwargs = kwargs + self.closed = False + self.instances.append(self) + + def close(self) -> None: + self.closed = True + + +class _FakeChannel: + instances: list["_FakeChannel"] = [] + + def __init__(self, consumer_world_size: int, *, timeout: int) -> None: + self.consumer_world_size = consumer_world_size + self.timeout = timeout + self.closed = False + self.instances.append(self) + + def close(self) -> None: + self.closed = True + + +def _parallel_runtime_config() -> ModelRuntimeConfig: + return ModelRuntimeConfig( + parallel_config=ParallelConfig(device_ids=[0, 1], sp_ulysses_degree=2, timeout=123), + ) + + +def test_parallel_vae_and_denoise_use_bidirectional_tensor_channels(monkeypatch) -> None: + _FakeWorker.instances = [] + _FakeChannel.instances = [] + monkeypatch.setattr(wan22_video, "VAEStage", _FakeStage) + monkeypatch.setattr(wan22_video, "MoeDitDenoisingStage", _FakeStage) + monkeypatch.setattr(wan22_video, "TextEncodingStage", _FakeStage) + monkeypatch.setattr(wan22_video, "ParallelWorker", _FakeWorker) + monkeypatch.setattr(wan22_video, "WorkerTensorChannel", _FakeChannel) + module_manager = Mock() + module_manager.get_model_info.return_value = {} + config = wan22_video.Wan22VideoPipelineConfig( + vae_config=_parallel_runtime_config(), + dit_high_config=_parallel_runtime_config(), + dit_low_config=_parallel_runtime_config(), + enable_vae_parallel=True, + enable_denoising_parallel=True, + ) + pipeline = wan22_video.Wan22VideoPipeline(device="cuda", torch_dtype=torch.bfloat16) + + pipeline.init(module_manager, config) + + assert pipeline._uses_direct_tensor_handoff + assert len(_FakeChannel.instances) == 2 + vae_to_denoise, denoise_to_vae = _FakeChannel.instances + vae_worker, denoise_worker = _FakeWorker.instances + assert vae_worker.kwargs == { + "tensor_output_channel": vae_to_denoise, + "tensor_output_methods": ("process",), + "tensor_input_channels": (denoise_to_vae,), + } + assert denoise_worker.kwargs == { + "tensor_output_channel": denoise_to_vae, + "tensor_output_methods": ("process",), + "tensor_input_channels": (vae_to_denoise,), + } + + pipeline.close() + + assert all(worker.closed for worker in _FakeWorker.instances) + assert all(channel.closed for channel in _FakeChannel.instances) diff --git a/tests/unit/test_run_examples.py b/tests/unit/test_run_examples.py index e5ed124..39713c7 100644 --- a/tests/unit/test_run_examples.py +++ b/tests/unit/test_run_examples.py @@ -129,6 +129,17 @@ def get_pipeline(parallelism: int, expert_backend: str, refiner_batch_cfg: bool) ) == (4, "sorted", True) +def test_prepare_sdpa_regression_forces_sdpa_and_disables_xformers(monkeypatch: pytest.MonkeyPatch) -> None: + from diffusers.utils import import_utils + + monkeypatch.setattr(import_utils, "_xformers_available", True) + + overrides = run_examples._prepare_sdpa_regression({"attn_impl": "FLASH_ATTN_4", "compile": True}) + + assert overrides == {"attn_impl": "TORCH_SDPA", "compile": True} + assert not import_utils._xformers_available + + def test_call_run_preserves_missing_negative_prompt_default() -> None: module = ModuleType("negative_prompt_example") From 540b5798d586d6a69ef924b79e4dcb21e1583472 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Mon, 3 Aug 2026 09:03:27 +0000 Subject: [PATCH 09/11] perf(distributed): unify high-throughput tensor communication Centralize process-group collectives and submit independent communication before synchronization across sequence-parallel model paths. Harden worker tensor channels with pooled CUDA IPC buffers, rank-local sharding, generation acknowledgements, cancellation cleanup, and resident LingBot condition/latent handoff. Document the communication architecture in English and Chinese, and add focused unit, integration, and SGLang comparison coverage. Verification:\n- ruff check on changed Python files\n- pytest focused communication and LingBot suites (58 passed)\n- mkdocs build --strict\n- LingBot World V2 four-GPU 77-frame benchmark and example --- README.md | 1 + README_zh.md | 1 + docs/en/communication.md | 309 ++++++++++++ docs/en/index.md | 1 + docs/en/parallel.md | 40 +- docs/zh/communication.md | 295 +++++++++++ docs/zh/index.md | 1 + docs/zh/parallel.md | 38 +- mkdocs.yml | 2 + telefuser/distributed/collectives.py | 74 +++ telefuser/distributed/parallel_shard.py | 16 +- telefuser/distributed/ulysses_comm.py | 15 - telefuser/distributed/vae_spatial.py | 29 +- telefuser/models/lingbot_video_dit.py | 42 +- telefuser/models/liveact_dit.py | 9 +- telefuser/models/video_projector.py | 5 +- telefuser/models/wan22_video_vae.py | 19 +- telefuser/models/wan_video_vae.py | 35 +- .../pipelines/lingbot_world_fast/pipeline.py | 10 +- .../pipelines/lingbot_world_fast/streaming.py | 252 +++++++++- .../pipelines/lingbot_world_fast/vae_stage.py | 2 + telefuser/worker/parallel_worker.py | 25 + telefuser/worker/ray_worker.py | 15 +- telefuser/worker/tensor_channel.py | 461 ++++++++++++++++-- tests/integration/test_collectives.py | 52 ++ .../integration/test_wan_video_vae_spatial.py | 99 ++++ .../integration/test_worker_tensor_channel.py | 106 ++++ tests/unit/distributed/test_collectives.py | 61 +++ tests/unit/distributed/test_parallel_shard.py | 38 +- tests/unit/models/test_lingbot_video_dit.py | 45 ++ .../lingbot_world_fast/test_module_loading.py | 2 +- .../lingbot_world_fast/test_streaming.py | 87 +++- tests/unit/worker/test_ray_worker.py | 48 ++ tests/unit/worker/test_tensor_channel.py | 86 +++- .../benchmark_tensor_channel_vs_sglang.py | 283 +++++++++++ 35 files changed, 2387 insertions(+), 217 deletions(-) create mode 100644 docs/en/communication.md create mode 100644 docs/zh/communication.md create mode 100644 telefuser/distributed/collectives.py create mode 100644 tests/integration/test_collectives.py create mode 100644 tests/integration/test_wan_video_vae_spatial.py create mode 100644 tests/unit/distributed/test_collectives.py create mode 100644 tests/unit/worker/test_ray_worker.py create mode 100644 tools/validation/benchmark_tensor_channel_vs_sglang.py diff --git a/README.md b/README.md index c769c37..3ebebdf 100644 --- a/README.md +++ b/README.md @@ -240,6 +240,7 @@ See [examples/README.md](examples/README.md) for the example runner and baseline - [docs/en/stream_server.md](docs/en/stream_server.md): LiveKit streaming, session APIs, data topics, and deployment - [docs/en/stream_scheduler.md](docs/en/stream_scheduler.md): actor-based stage scheduling, backpressure, lifecycle, metrics, and LingBot placement - [docs/en/parallel.md](docs/en/parallel.md): distributed inference architecture +- [docs/en/communication.md](docs/en/communication.md): collectives, CUDA IPC, synchronization, and transport lifecycle - [docs/en/latent_cache.md](docs/en/latent_cache.md): CacheSeek latent cache integration - [docs/en/feature_cache.md](docs/en/feature_cache.md): `AdaTaylorCache` - [docs/en/model_loading.md](docs/en/model_loading.md): model loading patterns diff --git a/README_zh.md b/README_zh.md index 2dfa87a..e88d415 100644 --- a/README_zh.md +++ b/README_zh.md @@ -232,6 +232,7 @@ telefuser/ - [docs/zh/stream_server.md](docs/zh/stream_server.md)：LiveKit 流服务、session API、data topic 和部署 - [docs/zh/stream_scheduler.md](docs/zh/stream_scheduler.md)：基于 actor 的 Stage 调度、backpressure、生命周期、指标和 LingBot 卡位 - [docs/zh/parallel.md](docs/zh/parallel.md)：分布式推理架构 +- [docs/zh/communication.md](docs/zh/communication.md)：collective、CUDA IPC、同步与传输生命周期 - [docs/zh/latent_cache.md](docs/zh/latent_cache.md)：CacheSeek latent cache 集成 - [docs/zh/feature_cache.md](docs/zh/feature_cache.md)：`AdaTaylorCache` - [docs/zh/model_loading.md](docs/zh/model_loading.md)：模型加载方式 diff --git a/docs/en/communication.md b/docs/en/communication.md new file mode 100644 index 0000000..4a7886d --- /dev/null +++ b/docs/en/communication.md @@ -0,0 +1,309 @@ +# Communication Architecture + +This guide describes how TeleFuser moves tensors between GPUs and worker processes, how communication ownership is +split across modules, and which invariants keep the implementation correct and efficient. For configuration of DP, +CFG, SP, PP, TP, and FSDP, see the [Parallel Inference Guide](parallel.md). + +## Design Goals + +TeleFuser communication follows five rules: + +1. Keep the control plane separate from large tensor data. +2. Keep reusable collective mechanics out of model implementations. +3. Avoid host staging and device-wide synchronization on same-host GPU paths. +4. Bound retained GPU memory and provide explicit cancellation and shutdown semantics. +5. Preserve a native PyTorch fallback when a specialized transport is not applicable. + +The implementation is layered rather than hidden behind one universal transport. NCCL collectives, CUDA IPC, Ray, +and service networking have different topology and lifetime requirements, so they share ownership rules but not one +runtime protocol. + +## Architecture Overview + +```text +Pipeline / model code + | declares topology, tensor layout, and stage connections + v +Parallel strategy and worker adapters + | DeviceMesh groups, Ulysses, Ring, VAE spatial, ParallelWorker + v +Shared communication mechanisms + | collectives.py | worker/tensor_channel.py + | process-group tensor movement | same-host cross-worker tensor movement + v v +PyTorch distributed / NCCL CUDA IPC + multiprocessing metadata +``` + +The main ownership boundaries are: + +| Area | Owner | Responsibility | +|------|-------|----------------| +| Process topology | `telefuser/distributed/device_mesh.py` | Build and expose named process groups | +| Shared collectives | `telefuser/distributed/collectives.py` | Contiguous gather buffers and grouped reductions | +| Sequence attention | `ulysses_comm.py`, `ring.py` | Strategy-specific All-to-All and P2P protocols | +| Sequence/CFG shards | `parallel_shard.py` | Tensor padding, slicing, gathering, and restoration | +| Spatial VAE | `vae_spatial.py` | Height shards, neighbor halos, and output gathering | +| Pipeline P2P | `pp_comm.py` | Rank-to-rank communication inside a PP process group | +| Worker execution | `worker/parallel_worker.py` | Process-group and worker lifecycle, command dispatch | +| Cross-worker tensors | `worker/tensor_channel.py` | CPU shared memory and persistent CUDA IPC pools | +| Cluster actors | `worker/ray_worker.py` | Ray resource assignment and optional local workers | + +Model code owns model-specific layout and reconstruction. It should call a shared collective or a strategy module +instead of allocating rank buffers or invoking tensor collectives directly. + +## Process Groups and Device Mesh + +`ParallelWorker` starts one spawned process per local rank. For groups larger than one rank it selects the assigned +device, initializes the platform distributed backend, and asks the stage to parallelize its models. CUDA normally +uses NCCL through the platform abstraction. + +`create_device_mesh_from_config()` creates named dimensions in this order: + +```text +DP -> CFG -> SP (ring, ulysses) -> PP -> TP +``` + +When Ring and Ulysses are both enabled, SP is a two-dimensional `(ring, ulysses)` submesh. Accessors such as +`get_cfg_group()`, `get_ring_group()`, `get_ulysses_group()`, and `get_pp_group()` prevent models from reconstructing +rank lists independently. The configured world size must equal the product of all enabled degrees. SP and TP are +currently mutually exclusive. + +## Shared Collective Primitives + +`telefuser/distributed/collectives.py` is an internal implementation boundary. It is intentionally not exported from +the top-level `telefuser.distributed` API. + +### Equal-shape gather + +`all_gather_stacked()` gathers equal-shaped local tensors into one rank-major allocation: + +```text +local [D0, ...] + -> all_gather_into_tensor +buffer [world_size * D0, ...] + -> view +result [world_size, D0, ...] +``` + +This replaces one allocation per rank with one contiguous output buffer. Consumers can retain the rank dimension for +model-specific reconstruction or use `all_gather_cat()` to concatenate along any tensor dimension in rank order. + +`parallel_shard.py`, LingBot Video sequence restoration, Wan/Wan2.2 VAE reconstruction, and VAE spatial gathering use +these primitives. Unequal VAE height shards are padded to the maximum local height before gather and trimmed after it. + +### Grouped reductions + +`all_reduce_sum_()` submits all independent sum reductions asynchronously before waiting for their work handles. Tile +blending uses it for value and weight tensors, avoiding duplicated synchronization code in model implementations. + +These helpers assume that all participating ranks call collectives in the same order with compatible shapes and +dtypes. Violating collective order is a distributed deadlock, not a recoverable per-rank error. + +## Sequence Parallel Communication + +### Ulysses + +Ulysses converts sequence shards with global heads into full-sequence tensors with local heads: + +```text +[B, S_local, H_global, D] + -> All-to-All scatter heads / gather sequence +[B, S_global, H_local, D] + -> local attention + -> All-to-All gather heads / scatter sequence +[B, S_local, H_global, D] +``` + +`ulysses_scatter_heads()` and `ulysses_gather_heads()` use functional `all_to_all_single` collectives. They return wait +closures so callers can separate submission from consumption. Attention implementations submit Q, K, and V before +waiting for any of them, allowing NCCL to schedule the three transfers without Python-side serialization. The output +All-to-All restores the original sequence/head layout. + +The number of attention heads must be divisible by the Ulysses world size, and the gathered sequence length must be +divisible on the inverse path. The helpers validate these constraints before communication. + +### Ring Attention + +Ring Attention keeps Q local and rotates K/V through neighboring ranks. `RingP2PComm` resolves group-local neighbors +to global ranks, batches `isend` and `irecv` operations with `batch_isend_irecv`, and reuses caller-provided receive +buffers when available. + +K and V may be concatenated into one transfer and split into views after receive. Communication for the next KV block +is submitted before attention on the current block; the implementation waits only before consuming the next block. +Partial attention results are combined with an online log-sum-exp merge. + +The AllGather Ring variant is simpler but materializes global K/V on every rank. It is a memory-heavy alternative, not +the preferred long-context path. + +## Spatial VAE Communication + +Height-sharded VAE decode has two communication patterns: + +1. Neighbor halo exchange before a spatial convolution. +2. Rank-ordered gather when a full-height tensor is required. + +Halo exchange uses reusable send and receive buffers and one `batch_isend_irecv` call for the available top and bottom +neighbors. Boundary ranks fill missing halos with zero. Buffer reuse avoids allocating contiguous halo tensors at +every layer and every frame. + +Full-height reconstruction uses the shared contiguous gather primitive. Rank-local heights are gathered first so +uneven shards can be padded and trimmed correctly. The final tensor restores its original channels-last or contiguous +memory format. + +## Cross-Worker Tensor Channel + +`WorkerTensorChannel` connects one producer worker group to one consumer worker group on the same host. It separates +small control metadata from tensor storage: + +```text +Producer worker Parent / control path Consumer worker + | | | + | stage tensor | | + |-- stage into IPC slot ---------->| WorkerTensorRef metadata ---->| + | | |-- map pool once + |<------- generation ACK / completion event -----------------------|-- peer copy +``` + +The parent receives `WorkerTensorRef` objects and never materializes CUDA tensor contents. Nested dictionaries, +tuples, and lists preserve their structure; duplicate tensor leaves are transported once per artifact. + +### Persistent CUDA IPC pools + +Stable CUDA tensor profiles are keyed by tensor index, shape, dtype, and source device. Each profile owns a persistent +allocation with two slots by default. Slots are selected round-robin so sequential traffic uses real double buffering. + +The pool allocation and IPC handle are created once. Consumers cache imported storage and event handles, so steady +state does not reopen CUDA IPC allocations. At most eight profiles are pooled per channel. Additional dynamic profiles +fall back to PyTorch multiprocessing tensor transport rather than retaining unbounded HBM. + +When `shard_dim` is set, every consumer rank receives only its rank-local view. The producer stages the tensor once, +and aggregate peer-copy traffic remains one logical tensor rather than one full tensor per consumer. LingBot's spatial +VAE uses height sharding with `shard_dim=-2`. + +### Stream ordering protocol + +Each slot has a reusable producer-ready event: + +1. The producer copies the source tensor into the slot on its current stream. +2. The producer records the ready event and publishes its handle with metadata. +3. The consumer stream waits for the ready event only when it is not already complete. +4. The consumer copies its mapped slot view into an output tensor. +5. The consumer records a reusable completion event before publishing its generation acknowledgement. +6. Before overwriting a reused slot, the producer stream waits for completion events from ranks that copied it. + +This protocol contains no device-wide synchronization in the transport path. Event `query()` provides a fast path +when producer staging or consumer copy has already completed. + +Acknowledgements use a lock-free shared generation array. A positive generation means that a rank copied the payload; +a negative generation means that it discarded the payload. The producer waits for completion metadata only from +ranks that performed a copy, so cancellation cannot wait for an event that was never recorded. + +### CPU and fallback transport + +CPU tensors use multiprocessing shared memory. CUDA profiles that cannot be pooled use PyTorch's multiprocessing CUDA +tensor transport. In both cases, one FIFO exists per consumer rank, and the receiving process performs final device +placement. + +## Control Plane and Lifecycle + +`ParallelWorker` command and result queues carry method names, arguments, small results, and tensor references. They +use `SimpleQueue` to avoid background feeder scheduling tails. Large tensors connected through a +`WorkerTensorChannel` stay on the direct data path. + +The channel contract is ordered and bounded: + +- Exactly one producer and one consumer group bind to a channel. +- Consumer rank count must match the channel configuration. +- Consumers resolve artifacts in producer order. +- A terminal cancelled artifact must be released with `discard_tensor_refs(..., sync=True)`. +- Shutdown stops consumers before producers and closes the channel last. +- Worker cleanup synchronizes pending device work before releasing local IPC mappings. + +Timeouts mark a worker failed and terminate its processes. Reusing a failed worker is rejected rather than risking a +partially ordered channel. + +## Pipeline Parallel and Ray Boundaries + +`PipelineP2PComm` is a different transport from `WorkerTensorChannel`. It communicates between ranks inside one PP +process group with NCCL send/recv and batched P2P operations. Existing Wan PP shape/grid broadcasts and latent +convenience methods remain owned by that PP path. + +CUDA IPC is a same-host mechanism. `RayWorker` respects the logical devices assigned by Ray and may host a local +`ParallelWorker`, but TeleFuser does not replace Ray's cross-node object transport with CUDA IPC. A deployment that +needs cross-node GPU-direct transfer requires a separately designed transport and topology contract. + +## Efficiency Invariants + +The communication implementation preserves these performance properties: + +- No parent-process CUDA tensor materialization on direct worker edges. +- No host staging on the pooled same-host CUDA path. +- Two logical device copies for a full handoff: producer staging and consumer output copy. +- Persistent pool, storage, and event handles in steady state. +- Bounded slots and profile count, preventing unbounded retained HBM. +- Stream events instead of device-wide synchronization. +- Rank-local copying when consumers operate on disjoint shards. +- One contiguous output allocation for equal-shape gather. +- Q/K/V collective submission before waits in Ulysses. +- Batched neighbor P2P and reusable halo buffers in spatial VAE and Ring paths. + +## Verification and Benchmarking + +Focused tests cover pure collective layout, real two-GPU NCCL ordering, CUDA IPC readiness and slot reuse, cancellation, +multi-consumer acknowledgement, and spatial VAE parity: + +```bash +pytest tests/unit/distributed/ +pytest tests/integration/test_collectives.py +pytest tests/integration/test_worker_tensor_channel.py +pytest tests/integration/test_wan_video_vae_spatial.py +``` + +The local SGLang comparison includes producer staging, metadata transport, target copy, target synchronization, and +slot acknowledgement for both implementations: + +```bash +python tools/validation/benchmark_tensor_channel_vs_sglang.py +``` + +Its default gate uses 200 measured transfers. TeleFuser p50 must be no more than 5% above SGLang, while p95 allows the +larger of 10% or 0.05 ms to account for sub-millisecond multiprocessing scheduling jitter. Compare copy counts and +mean latency as well as percentiles; a single process-scheduling tail is not evidence of a transport regression. + +Pipeline-level validation should rerun every pipeline whose communication call site changed. The example runner +provides baseline output comparisons: + +```bash +python examples/run_examples.py --pipeline --gpus 0,1,2,3 +``` + +## Extension Rules + +When adding a communication path: + +1. Put generic equal-shape gather or reduction mechanics in `distributed/collectives.py`. +2. Put algorithm-specific protocols in a focused module under `telefuser/distributed/`. +3. Keep model code responsible only for tensor layout and model semantics. +4. Use `WorkerTensorChannel` only for a same-host, single-producer/single-consumer-group edge. +5. Do not add a new fallback, environment variable, or public configuration field without a concrete topology gap. +6. Specify ordering, ownership, cancellation, timeout, and shutdown before optimizing the happy path. +7. Add a real multi-process test for any new collective or IPC synchronization rule. + +Do not route large tensors through the parent merely because the control path already exists. Do not use a device-wide +synchronize to repair an ordering bug; express the dependency with process-group work handles or stream events. + +## Known Limits + +- CUDA IPC pools are same-host only. +- Stable pooled profiles require fixed tensor index, shape, dtype, and source device. +- Ring AllGather trades implementation simplicity for replicated K/V memory. +- Spatial VAE halo exchange reuses buffers but currently waits before the dependent convolution. +- WAN pipeline-parallel communication remains a separate, model-specific compatibility area. +- Ray cross-node tensor performance depends on Ray transport and cluster configuration. + +## Related Documentation + +- [Parallel Inference Guide](parallel.md) +- [Attention Implementation Guide](attention.md) +- [Streaming Scheduler](stream_scheduler.md) +- [Testing Guide](testing.md) diff --git a/docs/en/index.md b/docs/en/index.md index affd8a2..acee67a 100644 --- a/docs/en/index.md +++ b/docs/en/index.md @@ -105,6 +105,7 @@ telefuser stream-serve examples/lingbot/lingbot_world_fast_image_to_video_h100.p ConfigurationRuntime, attention, quantization, and offload settings. TF-KernelInstall, build, verify, and use the optional CUDA extension. Parallel InferenceDistributed processing strategies. +Communication ArchitectureNCCL collectives, CUDA IPC, ordering, and efficiency. Adding New ModelIntegrate new model architectures and stages. ProfilerPerformance analysis tools. diff --git a/docs/en/parallel.md b/docs/en/parallel.md index 4129711..01b228d 100644 --- a/docs/en/parallel.md +++ b/docs/en/parallel.md @@ -1,6 +1,8 @@ # Parallel Inference Guide This document provides a detailed introduction to TeleFuser's distributed parallel inference architecture, including principles, configuration methods, and usage examples. +For tensor dataflow, synchronization, transport ownership, and performance invariants, see the +[Communication Architecture](communication.md). ## Overview @@ -46,14 +48,19 @@ device_mesh = create_device_mesh_from_config(config) ``` telefuser/distributed/ ├── device_mesh.py # DeviceMesh creation and process group management +├── collectives.py # Shared contiguous gather and reduction primitives ├── pp_comm.py # Pipeline parallel P2P communication ├── ulysses_comm.py # Ulysses All-to-All communication primitives ├── ring.py # Ring Attention P2P communication ├── parallel_shard.py # Sequence parallel tensor shard/unshard +├── vae_spatial.py # Height-sharded VAE halo exchange ├── fsdp.py # FSDP data parallel └── tp_parallelize.py # Tensor parallel utilities ``` +Model code owns tensor layout and reconstruction semantics, while reusable collective buffer allocation and reduction +submission stay in `collectives.py`. Strategy-specific protocols remain in their corresponding modules. + ## Sequence Parallelism Sequence parallelism is used to process very long sequences (e.g., long videos) by splitting the sequence dimension across multiple GPUs. @@ -161,10 +168,13 @@ Split model layers across multiple GPUs for large model inference. ### Cross-worker tensor channels -Independent `ParallelWorker` groups can connect adjacent stages with a `WorkerTensorChannel`. The producer sends -tensor storage directly to every consumer rank through multiprocessing shared memory or CUDA IPC. The parent process -receives only a `WorkerTensorRef` containing the channel, transfer, shape, dtype, and source-device metadata. The -consumer resolves that reference on its own device before invoking the unchanged stage method. +Independent `ParallelWorker` groups can connect adjacent stages with a `WorkerTensorChannel`. CPU tensors use +multiprocessing shared memory. CUDA tensors use two bounded producer-owned IPC slots per stable tensor profile; each +IPC allocation is opened once and then reused. Pool handles and rank-local offsets travel as private +`WorkerTensorRef` metadata through the existing control path. Reusable interprocess CUDA events order producer +staging before consumer copies and order the next slot write after every consumer copy. Consumers record completion +events before publishing generation acknowledgements; the producer waits on those events in its staging stream +without device-wide synchronization. ```python from telefuser.worker import ParallelWorker, WorkerTensorChannel @@ -178,15 +188,33 @@ denoise_worker = ParallelWorker( vae_worker = ParallelWorker(vae_stage, tensor_input_channels=(latent_channel,)) ``` +Set `shard_dim` when consumer ranks operate on disjoint tensor slices. The producer stages the tensor once and each +rank copies only its slice. LingBot's spatial VAE uses `shard_dim=-2`, so aggregate peer-copy traffic stays at one +logical latent instead of growing with the VAE world size. Including producer-local staging, the device-copy budget is +two logical latents. At most eight stable CUDA tensor profiles are pooled per channel; additional dynamic profiles +fall back to PyTorch CUDA IPC instead of growing retained HBM without bound. + This is a point-to-point, single-producer/single-consumer-group path. Enable it only for outputs whose complete consumer set is the connected worker group. Calls that need to inspect a tensor in the parent may pass -`_tensor_transport=False`. Start both workers before submitting work, stop both workers before closing the channel, +`_tensor_transport=False`. Start both workers before submitting work, stop consumers before producers, then close the channel, preserve producer order at the consumer, and treat transported tensors as immutable until the consumer finishes. -The receiver discards older FIFO entries when the scheduler cancels an artifact before consumption. +Schedulers that cancel a terminal artifact must call the consumer worker's +`discard_tensor_refs(ref, sync=True)` in producer order. This releases CPU shared memory or CUDA IPC storage without +materializing the tensor in the parent. A normal receive also discards older cancelled FIFO entries as a fallback. Regular worker dispatch also sends shared-memory or CUDA IPC handles to each rank and lets the receiving rank perform the final device placement. The parent does not allocate a temporary copy on every target GPU. +Use the local SGLang checkout to run the end-to-end latency gate on the same GPUs and tensor shape: + +```bash +python tools/validation/benchmark_tensor_channel_vs_sglang.py +``` + +The comparison includes producer staging, metadata transport, target copy, target synchronization, and slot +acknowledgement for both implementations. Its default 200-sample gate rejects p50 regressions above 5% and p95 +regressions above the larger of 10% or 0.05 ms, accounting for sub-millisecond multiprocessing scheduling jitter. + ### Principle ``` diff --git a/docs/zh/communication.md b/docs/zh/communication.md new file mode 100644 index 0000000..e13c34f --- /dev/null +++ b/docs/zh/communication.md @@ -0,0 +1,295 @@ +# 通信架构 + +本文说明 TeleFuser 如何在 GPU、进程和 worker 之间传递 tensor，通信职责如何划分，以及实现通过哪些约束 +保证正确性和效率。DP、CFG、SP、PP、TP 与 FSDP 的配置方法见[并行推理指南](parallel.md)。 + +## 设计目标 + +TeleFuser 的通信实现遵循五条原则： + +1. 控制面与大 tensor 数据面分离。 +2. 可复用的 collective 机制不放在模型实现中。 +3. 同机 GPU 路径避免 host staging 和 device-wide synchronization。 +4. 限制常驻显存，并明确取消、超时和关闭语义。 +5. 优化传输不适用时保留 PyTorch 原生回退路径。 + +实现采用分层设计，而不是用一个通用 transport 隐藏所有差异。NCCL collective、CUDA IPC、Ray 和服务网络 +具有不同的拓扑与生命周期要求；它们共享职责规则，但不共享同一种运行时协议。 + +## 架构总览 + +```text +Pipeline / 模型代码 + | 声明拓扑、tensor layout 和 stage 连接关系 + v +并行策略与 worker adapter + | DeviceMesh group、Ulysses、Ring、VAE spatial、ParallelWorker + v +共享通信机制 + | collectives.py | worker/tensor_channel.py + | 进程组内 tensor 搬运 | 同机跨 worker tensor 搬运 + v v +PyTorch distributed / NCCL CUDA IPC + multiprocessing 元数据 +``` + +主要职责边界如下： + +| 范围 | 负责模块 | 职责 | +|------|----------|------| +| 进程拓扑 | `telefuser/distributed/device_mesh.py` | 创建并暴露具名 process group | +| 共享 collective | `telefuser/distributed/collectives.py` | 连续 gather buffer 与成组 reduction | +| 序列注意力 | `ulysses_comm.py`、`ring.py` | 策略专用的 All-to-All 与 P2P 协议 | +| 序列/CFG 分片 | `parallel_shard.py` | tensor padding、切分、gather 和恢复 | +| 空间 VAE | `vae_spatial.py` | 高度分片、邻居 halo 和输出 gather | +| Pipeline P2P | `pp_comm.py` | PP process group 内 rank-to-rank 通信 | +| Worker 执行 | `worker/parallel_worker.py` | 进程组、worker 生命周期和命令派发 | +| 跨 worker tensor | `worker/tensor_channel.py` | CPU shared memory 和持久 CUDA IPC pool | +| 集群 actor | `worker/ray_worker.py` | Ray 资源分配和可选的本地 worker group | + +模型代码负责模型特有的 layout 和重建语义；不应自行分配 rank buffer 或直接调用 tensor collective，而应 +调用共享 collective 或策略模块。 + +## 进程组与 DeviceMesh + +`ParallelWorker` 为每个本地 rank 启动一个 spawn 进程。当 group 大于一个 rank 时，它选择分配的 device， +初始化平台对应的 distributed backend，再让 stage 并行化模型。CUDA 平台通常通过平台抽象使用 NCCL。 + +`create_device_mesh_from_config()` 按以下顺序创建具名维度： + +```text +DP -> CFG -> SP (ring, ulysses) -> PP -> TP +``` + +同时启用 Ring 和 Ulysses 时，SP 是二维 `(ring, ulysses)` 子 mesh。`get_cfg_group()`、 +`get_ring_group()`、`get_ulysses_group()` 和 `get_pp_group()` 等 accessor 避免模型重复推导 rank 列表。 +配置的 world size 必须等于所有并行 degree 的乘积。目前 SP 和 TP 互斥。 + +## 共享 Collective 原语 + +`telefuser/distributed/collectives.py` 是内部实现边界，有意不从顶层 `telefuser.distributed` API 导出。 + +### 等形状 Gather + +`all_gather_stacked()` 把各 rank 的等形状 tensor gather 到一个 rank-major allocation： + +```text +local [D0, ...] + -> all_gather_into_tensor +buffer [world_size * D0, ...] + -> view +result [world_size, D0, ...] +``` + +这样每次 gather 只分配一个连续输出 buffer，不再为每个 rank 单独分配 tensor。调用方可以保留 rank 维进行 +模型特有的重建，也可以使用 `all_gather_cat()` 按 rank 顺序沿任意维拼接。 + +`parallel_shard.py`、LingBot Video 序列恢复、Wan/Wan2.2 VAE 重建和 VAE spatial gather 都复用这些原语。 +VAE 高度分片不等长时，先 pad 到最大本地高度，gather 后再裁剪。 + +### 成组 Reduction + +`all_reduce_sum_()` 先异步提交所有相互独立的 sum reduction，再统一等待 work handle。Tile blending 用它同时 +归约 value 和 weight，模型中不再重复实现同步逻辑。 + +这些 helper 要求所有参与 rank 以相同顺序调用 collective，并提供兼容的 shape 和 dtype。Collective 顺序 +不一致会造成分布式死锁，不能作为单 rank 异常恢复。 + +## 序列并行通信 + +### Ulysses + +Ulysses 把“序列分片、完整 heads”转换成“完整序列、本地 heads”： + +```text +[B, S_local, H_global, D] + -> All-to-All：scatter heads / gather sequence +[B, S_global, H_local, D] + -> 本地 attention + -> All-to-All：gather heads / scatter sequence +[B, S_local, H_global, D] +``` + +`ulysses_scatter_heads()` 和 `ulysses_gather_heads()` 使用 functional `all_to_all_single`，并返回 wait closure， +使提交与消费分离。Attention 会先提交 Q、K、V 三个 collective，再等待其中任何一个，让 NCCL 不受 Python +串行提交限制。输出 All-to-All 恢复原始 sequence/head layout。 + +Attention head 数必须能被 Ulysses world size 整除，反向 layout 恢复时 gathered sequence length 也必须可整除。 +Helper 会在通信前验证这些约束。 + +### Ring Attention + +Ring Attention 保持 Q 本地不动，让 K/V 在相邻 rank 之间轮转。`RingP2PComm` 把 group-local 邻居解析为 +global rank，使用 `batch_isend_irecv` 批量提交 `isend` 和 `irecv`，并优先复用调用方提供的 receive buffer。 + +K/V 可以拼成一次传输，接收后再切成 view。下一块 KV 的通信先提交，再计算当前块 attention，只在消费下一块 +之前等待。各块的 attention 结果通过 online log-sum-exp 合并。 + +Ring AllGather 变体实现更简单，但每个 rank 都会物化全局 K/V；它是显存开销更大的备选实现，不是长上下文 +首选路径。 + +## 空间 VAE 通信 + +按高度分片的 VAE decode 包含两类通信： + +1. 空间卷积前的邻居 halo 交换。 +2. 需要完整高度 tensor 时的 rank-order gather。 + +Halo 交换复用 send/receive buffer，并用一次 `batch_isend_irecv` 提交当前 rank 存在的上下邻居操作。边界 rank +把缺失 halo 填零。Buffer 复用避免每层、每帧重复创建 contiguous halo tensor。 + +完整高度重建使用共享连续 gather 原语。先 gather 各 rank 的本地高度，才能正确处理不等分片的 padding 和裁剪。 +最终 tensor 会恢复原来的 channels-last 或 contiguous memory format。 + +## 跨 Worker Tensor Channel + +`WorkerTensorChannel` 在同一主机上连接一个 producer worker group 和一个 consumer worker group。小型控制 +元数据与 tensor storage 分开传输： + +```text +Producer worker 父进程 / 控制路径 Consumer worker + | | | + | stage tensor | | + |-- staging 到 IPC slot ---------->| WorkerTensorRef 元数据 ------>| + | | |-- pool 只映射一次 + |<------- generation ACK / completion event -----------------------|-- peer copy +``` + +父进程只接收 `WorkerTensorRef`，不会物化 CUDA tensor 内容。嵌套 dict、tuple 和 list 保持结构；同一 artifact +中的重复 tensor leaf 只传输一次。 + +### 持久 CUDA IPC Pool + +稳定 CUDA tensor profile 由 tensor index、shape、dtype 和 source device 共同标识。每个 profile 持有一块 +持久 allocation，默认包含两个 slot。Slot 使用 round-robin 选择，使顺序流量真正使用双缓冲。 + +Pool allocation 和 IPC handle 只创建一次。Consumer 缓存导入后的 storage 和 event handle，steady state 不会 +反复打开 CUDA IPC allocation。每个 channel 最多缓存八个 profile；更多动态 profile 回退到 PyTorch +multiprocessing tensor transport，避免常驻 HBM 无界增长。 + +配置 `shard_dim` 后，每个 consumer rank 只接收自己的 rank-local view。Producer 只 staging 一次，聚合 +peer-copy 流量保持为一个 logical tensor，而不是为每个 consumer 拷贝一份完整 tensor。LingBot 空间 VAE 使用 +`shard_dim=-2` 按高度分片。 + +### Stream 顺序协议 + +每个 slot 持有一个可复用的 producer-ready event： + +1. Producer 在当前 stream 把 source tensor copy 到 slot。 +2. Producer 记录 ready event，并随元数据发布 event handle。 +3. Ready event 尚未完成时，consumer stream 才等待它。 +4. Consumer 把映射后的 slot view copy 到 output tensor。 +5. Consumer 先记录可复用 completion event，再发布 generation ACK。 +6. Producer 覆盖复用 slot 前，其 staging stream 等待所有实际执行 copy 的 rank completion event。 + +整个 transport 路径不包含 device-wide synchronization。若 producer staging 或 consumer copy 已完成，event +`query()` 提供快速路径。 + +ACK 使用 lock-free shared generation array。正 generation 表示该 rank 已 copy payload；负 generation 表示 +该 rank 已 discard。Producer 只为真正 copy 的 rank 等待 completion metadata，因此取消路径不会等待一个 +从未记录的 event。 + +### CPU 与回退传输 + +CPU tensor 使用 multiprocessing shared memory。无法池化的 CUDA profile 使用 PyTorch multiprocessing 的 +CUDA tensor transport。两种情况都为每个 consumer rank 保留一个 FIFO，最终 device placement 由接收进程完成。 + +## 控制面与生命周期 + +`ParallelWorker` 的 command/result queue 传输方法名、参数、小型结果和 tensor reference。它们使用 +`SimpleQueue`，避免 background feeder 引入调度尾延迟。通过 `WorkerTensorChannel` 连接的大 tensor 留在直接 +数据路径。 + +Channel contract 是有序且有界的： + +- 一个 channel 只能绑定一个 producer 和一个 consumer group。 +- Consumer rank 数必须与 channel 配置匹配。 +- Consumer 必须按 producer 顺序解析 artifact。 +- 被取消的末端 artifact 必须调用 `discard_tensor_refs(..., sync=True)` 释放。 +- 关闭时先停止 consumer，再停止 producer，最后关闭 channel。 +- Worker cleanup 在释放本地 IPC mapping 前同步尚未完成的 device work。 + +Timeout 会把 worker 标记为 failed 并终止其进程。失败 worker 不允许继续复用，避免破坏 channel 的部分顺序。 + +## Pipeline Parallel 与 Ray 边界 + +`PipelineP2PComm` 与 `WorkerTensorChannel` 是不同传输。它在同一 PP process group 内使用 NCCL send/recv 和 +batched P2P。现有 Wan PP 的 shape/grid broadcast 与 latent convenience method 仍由 PP 路径负责。 + +CUDA IPC 只能用于同一主机。`RayWorker` 遵守 Ray 分配的逻辑 device，并可在 actor 内运行本地 +`ParallelWorker`，但 TeleFuser 不会用 CUDA IPC 替代 Ray 的跨节点 object transport。需要跨节点 GPU-direct +传输时，必须单独设计 transport 和拓扑 contract。 + +## 效率约束 + +通信实现保持以下性能属性： + +- 直接 worker edge 不在父进程物化 CUDA tensor。 +- 同机池化 CUDA 路径不经过 host staging。 +- 完整 handoff 只包含两个 logical device copy：producer staging 和 consumer output copy。 +- Steady state 复用 pool、storage 和 event handle。 +- Slot 和 profile 数量有界，避免常驻 HBM 无界增长。 +- 使用 stream event，不使用 device-wide synchronization。 +- Consumer 处理互斥 shard 时只 copy rank-local 数据。 +- 等形状 gather 只分配一个连续输出 buffer。 +- Ulysses 的 Q/K/V collective 先全部提交再等待。 +- Spatial VAE 和 Ring 使用批量邻居 P2P，并复用 halo/receive buffer。 + +## 验证与基准测试 + +专项测试覆盖纯 collective layout、真实双卡 NCCL 顺序、CUDA IPC readiness 与 slot reuse、取消、多 consumer +ACK，以及空间 VAE parity： + +```bash +pytest tests/unit/distributed/ +pytest tests/integration/test_collectives.py +pytest tests/integration/test_worker_tensor_channel.py +pytest tests/integration/test_wan_video_vae_spatial.py +``` + +本地 SGLang 对比同时包含两种实现的 producer staging、metadata transport、target copy、target synchronization +和 slot ACK： + +```bash +python tools/validation/benchmark_tensor_channel_vs_sglang.py +``` + +默认门禁测量 200 次传输。TeleFuser p50 不得比 SGLang 高 5% 以上；p95 上限取 10% 和 0.05 ms 中较宽者， +用于覆盖亚毫秒 multiprocessing 调度抖动。判断回归时还应比较 copy 次数和 mean latency；单个进程调度尾点 +不能单独证明 transport 退化。 + +修改通信调用点后，应复跑所有受影响 pipeline。Example runner 提供 baseline 输出对比： + +```bash +python examples/run_examples.py --pipeline --gpus 0,1,2,3 +``` + +## 扩展规则 + +新增通信路径时： + +1. 通用的等形状 gather 或 reduction 机制放在 `distributed/collectives.py`。 +2. 算法专用协议放在 `telefuser/distributed/` 下的聚焦模块中。 +3. 模型代码只负责 tensor layout 和模型语义。 +4. `WorkerTensorChannel` 只用于同机、单 producer、单 consumer group 的 edge。 +5. 没有明确拓扑缺口时，不新增 fallback、环境变量或公共配置字段。 +6. 优化 happy path 前，先定义顺序、所有权、取消、timeout 和 shutdown。 +7. 新 collective 或 IPC 同步规则必须添加真实多进程测试。 + +不要因为控制路径已经存在就让大 tensor 绕行父进程；也不要用 device-wide synchronize 修补顺序问题，应通过 +process-group work handle 或 stream event 表达依赖。 + +## 已知边界 + +- CUDA IPC pool 只支持同一主机。 +- 稳定池化 profile 要求 tensor index、shape、dtype 和 source device 固定。 +- Ring AllGather 用实现简单性换取每个 rank 的 K/V 副本显存。 +- 空间 VAE halo exchange 已复用 buffer，但当前仍会在依赖它的卷积前等待。 +- WAN pipeline-parallel 通信仍是独立的模型专用兼容区域。 +- Ray 跨节点 tensor 性能取决于 Ray transport 与集群配置。 + +## 相关文档 + +- [并行推理指南](parallel.md) +- [Attention 实现指南](attention.md) +- [流式调度器](stream_scheduler.md) +- [测试指南](testing.md) diff --git a/docs/zh/index.md b/docs/zh/index.md index a224e06..6b1c92b 100644 --- a/docs/zh/index.md +++ b/docs/zh/index.md @@ -104,6 +104,7 @@ telefuser stream-serve examples/lingbot/lingbot_world_fast_image_to_video_h100.p 配置运行时、注意力、量化和卸载配置。 TF-Kernel安装、编译、验证和使用可选 CUDA 扩展。 并行推理分布式处理策略。 +通信架构NCCL collective、CUDA IPC、顺序与效率设计。 新增模型集成新的模型架构和阶段。 性能分析性能分析工具。 diff --git a/docs/zh/parallel.md b/docs/zh/parallel.md index b3e7a20..e2a8b4f 100644 --- a/docs/zh/parallel.md +++ b/docs/zh/parallel.md @@ -1,6 +1,7 @@ # 并行推理指南 本文档详细介绍 TeleFuser 的分布式并行推理架构，包括原理介绍、配置方法和使用示例。 +Tensor 数据流、同步协议、transport 职责和性能约束见[通信架构](communication.md)。 ## 概述 @@ -46,14 +47,19 @@ device_mesh = create_device_mesh_from_config(config) ``` telefuser/distributed/ ├── device_mesh.py # DeviceMesh 创建和进程组管理 +├── collectives.py # 共享的连续 gather 与 reduction 原语 ├── pp_comm.py # 流水线并行 P2P 通信 ├── ulysses_comm.py # Ulysses All-to-All 通信原语 ├── ring.py # Ring Attention P2P 通信 ├── parallel_shard.py # 序列并行张量分片/反分片 +├── vae_spatial.py # 按高度分片的 VAE halo 交换 ├── fsdp.py # FSDP 数据并行 └── tp_parallelize.py # 张量并行工具 ``` +模型代码只负责张量布局和重建语义；可复用的 collective buffer 分配与 reduction 提交统一放在 +`collectives.py`，策略专用协议仍保留在各自模块中。 + ## 序列并行 序列并行用于处理超长序列（如长视频），将序列维度分割到多个 GPU。 @@ -161,10 +167,13 @@ out = out_wait() ### 跨 Worker Tensor 通道 -相邻 stage 如果属于不同 `ParallelWorker` group，可以使用 `WorkerTensorChannel` 连接。Producer 通过 -multiprocessing shared memory 或 CUDA IPC，把 tensor storage 直接发送给每个 consumer rank。主进程只接收 -`WorkerTensorRef`，其中包含 channel、transfer、shape、dtype、字节数和源设备等元数据。Consumer 在调用原有 -stage 方法之前，负责在自己的设备上解析引用；stage 的函数签名无需改变。 +相邻 stage 如果属于不同 `ParallelWorker` group，可以使用 `WorkerTensorChannel` 连接。CPU tensor 使用 +multiprocessing shared memory；CUDA tensor 对每个稳定 tensor profile 使用两个有界的 producer-owned IPC +slot，每个 IPC allocation 只打开一次并持续复用。Pool handle 与各 rank 的 offset 作为私有 +`WorkerTensorRef` 元数据沿现有控制路径传输。双向复用的跨进程 CUDA event 既保证 producer staging 完成后 +consumer 才 copy，也保证所有 consumer copy 完成后 producer 才覆盖 slot。Consumer 先记录 completion +event 再发布 generation ACK，producer staging stream 等待这些 event；整个过程不需要 device-wide +synchronization。 ```python from telefuser.worker import ParallelWorker, WorkerTensorChannel @@ -178,14 +187,31 @@ denoise_worker = ParallelWorker( vae_worker = ParallelWorker(vae_stage, tensor_input_channels=(latent_channel,)) ``` +当 consumer rank 只处理互不重叠的 tensor slice 时可设置 `shard_dim`。Producer 只 staging 一次，各 rank 仅 +copy 自己的 slice。LingBot 空间并行 VAE 使用 `shard_dim=-2`，因此跨卡 peer-copy 聚合通信量保持为一个 +logical latent，而不会随 VAE world size 放大；计入 producer 本地 staging 后，设备搬运预算为两个 logical +latent。每个 channel 最多池化八种稳定 CUDA tensor profile；更多动态 profile 会回退到 PyTorch CUDA IPC， +避免 retained HBM 无界增长。 + 该路径是单 producer、单 consumer group 的点对点 FIFO。只有 tensor 的完整 consumer 集合就是所连接的 worker group 时才能启用。需要在主进程读取 tensor 的调用可以传入 `_tensor_transport=False`。Consumer 必须保持 -producer 顺序，并在两个 worker 都停止后再关闭 channel。Scheduler 取消尚未消费的 artifact 时，receiver 会 -丢弃更早的 FIFO entry，避免污染后续传输。 +producer 顺序；关闭时先停止 consumer、再停止 producer，最后关闭 channel。Scheduler 取消末端 artifact 时，必须按 producer +顺序调用 consumer worker 的 `discard_tensor_refs(ref, sync=True)`。该操作不会在父进程物化 tensor，并会 +立即释放 CPU shared memory 或 CUDA IPC storage；正常 receive 仍会兜底丢弃更早的已取消 FIFO entry。 常规 worker 派发同样只向各 rank 发送 shared-memory 或 CUDA IPC handle，最终 device placement 由接收 rank 完成；主进程不再为每张目标 GPU 分配临时副本。 +可使用本地 SGLang checkout 在相同 GPU 和 tensor shape 上运行端到端延迟门禁： + +```bash +python tools/validation/benchmark_tensor_channel_vs_sglang.py +``` + +该比较同时计入两种实现的 producer staging、元数据传输、target copy、target 同步与 slot ACK。默认门禁使用 +200 个样本：p50 不得比 SGLang 高 5% 以上；p95 上限取 10% 与 0.05 ms 中较宽者，以覆盖亚毫秒级 +multiprocessing 调度抖动。 + ### 原理 ``` diff --git a/mkdocs.yml b/mkdocs.yml index 2673c3e..29fb4f4 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -126,6 +126,7 @@ plugins: User Guides: 用户指南 Benchmarks: 基准测试 Parallel Inference: 并行推理 + Communication Architecture: 通信架构 Developer Guides: 开发者指南 Configuration: 配置 Tools: 工具 @@ -178,6 +179,7 @@ nav: - Quantization: quantization.md - Parallel Inference: - parallel.md + - Communication Architecture: communication.md - attention.md - feature_cache.md - latent_cache.md diff --git a/telefuser/distributed/collectives.py b/telefuser/distributed/collectives.py new file mode 100644 index 0000000..9a5516e --- /dev/null +++ b/telefuser/distributed/collectives.py @@ -0,0 +1,74 @@ +"""Shared tensor collective primitives used by model parallel strategies.""" + +from __future__ import annotations + +from collections.abc import Iterable + +import torch +import torch.distributed as dist + + +def _resolve_world_size(group: dist.ProcessGroup | None, world_size: int | None) -> int: + if world_size is None: + world_size = dist.get_world_size(group=group) + if world_size < 1: + raise ValueError("world_size must be at least one") + return world_size + + +def all_gather_stacked( + tensor: torch.Tensor, + *, + group: dist.ProcessGroup | None = None, + world_size: int | None = None, +) -> torch.Tensor: + """Gather equal-shaped tensors into one rank-major contiguous buffer.""" + world_size = _resolve_world_size(group, world_size) + if world_size == 1: + return tensor.unsqueeze(0) + + gather_input = tensor.contiguous() + original_shape = gather_input.shape + if gather_input.ndim == 0: + gather_input = gather_input.reshape(1) + gathered = torch.empty( + (world_size * gather_input.shape[0], *gather_input.shape[1:]), + dtype=gather_input.dtype, + device=gather_input.device, + ) + dist.all_gather_into_tensor(gathered, gather_input, group=group) + return gathered.view(world_size, *original_shape) + + +def all_gather_cat( + tensor: torch.Tensor, + *, + dim: int, + group: dist.ProcessGroup | None = None, + world_size: int | None = None, +) -> torch.Tensor: + """Gather equal-shaped shards in rank order and concatenate one dimension.""" + if tensor.ndim == 0: + raise ValueError("all_gather_cat requires a tensor with at least one dimension") + dim = dim if dim >= 0 else tensor.ndim + dim + if not 0 <= dim < tensor.ndim: + raise ValueError(f"dim={dim} is invalid for a {tensor.ndim}D tensor") + world_size = _resolve_world_size(group, world_size) + if world_size == 1: + return tensor + + gather_input = tensor.movedim(dim, 0).contiguous() + gathered = all_gather_stacked(gather_input, group=group, world_size=world_size) + merged = gathered.flatten(0, 1).movedim(0, dim) + return merged.contiguous() + + +def all_reduce_sum_( + tensors: Iterable[torch.Tensor], + *, + group: dist.ProcessGroup | None = None, +) -> None: + """Sum tensors in place, submitting independent reductions before waiting.""" + works = [dist.all_reduce(tensor, group=group, async_op=True) for tensor in tensors] + for work in works: + work.wait() diff --git a/telefuser/distributed/parallel_shard.py b/telefuser/distributed/parallel_shard.py index 06325ed..6850f28 100644 --- a/telefuser/distributed/parallel_shard.py +++ b/telefuser/distributed/parallel_shard.py @@ -20,6 +20,7 @@ import torch.nn.functional as F from torch.distributed.device_mesh import DeviceMesh +from .collectives import all_gather_cat from .device_mesh import ( get_attention_strategy, get_cfg_group, @@ -189,10 +190,9 @@ def sequence_parallel_unshard( unshard_tensors = [] for tensor, seq_dim, seq_len in zip(tensors, seq_dims, seq_lens): - # All-gather across the shard group - unshard = [torch.zeros_like(tensor) for _ in range(shard_degree)] - dist.all_gather(unshard, tensor, group=shard_group) - unshard = torch.cat(unshard, dim=seq_dim).narrow(dim=seq_dim, start=0, length=seq_len) + unshard = all_gather_cat(tensor, dim=seq_dim, group=shard_group, world_size=shard_degree).narrow( + dim=seq_dim, start=0, length=seq_len + ) unshard_tensors.append(unshard) return unshard_tensors @@ -239,12 +239,8 @@ def cfg_parallel_unshard(device_mesh: DeviceMesh, tensors: list[torch.Tensor]) - unshard_tensors = [] for tensor in tensors: - unshard = torch.zeros( - (cfg_world_size, *tensor.shape[1:]), - dtype=tensor.dtype, - device=tensor.device, + unshard_tensors.append( + all_gather_cat(tensor, dim=0, group=get_cfg_group(device_mesh), world_size=cfg_world_size) ) - dist.all_gather_into_tensor(unshard, tensor, group=get_cfg_group(device_mesh)) - unshard_tensors.append(unshard) return unshard_tensors diff --git a/telefuser/distributed/ulysses_comm.py b/telefuser/distributed/ulysses_comm.py index 28f3643..faf5833 100644 --- a/telefuser/distributed/ulysses_comm.py +++ b/telefuser/distributed/ulysses_comm.py @@ -35,21 +35,6 @@ def _wait_async_tensor(tensor: torch.Tensor) -> torch.Tensor: return tensor -def ulysses_all_to_all_split_cat( - tensor: torch.Tensor, - process_group: dist.ProcessGroup, - *, - scatter_dim: int, - gather_dim: int, -) -> torch.Tensor: - """Run the source-style synchronous list all-to-all used by LingBot Ulysses.""" - _, world_size = _get_distributed_info(process_group) - inputs = [part.contiguous() for part in torch.tensor_split(tensor, world_size, scatter_dim)] - outputs = [torch.empty_like(inputs[0]) for _ in range(world_size)] - dist.all_to_all(outputs, inputs, group=process_group) - return torch.cat(outputs, dim=gather_dim).contiguous() - - def ulysses_scatter_heads( tensor: torch.Tensor, process_group: dist.ProcessGroup, diff --git a/telefuser/distributed/vae_spatial.py b/telefuser/distributed/vae_spatial.py index 2c4e2b3..85f2f41 100644 --- a/telefuser/distributed/vae_spatial.py +++ b/telefuser/distributed/vae_spatial.py @@ -7,6 +7,8 @@ import torch.nn as nn import torch.nn.functional as F +from .collectives import all_gather_stacked + def _spatial_rank() -> int: return dist.get_rank() if dist.is_initialized() else 0 @@ -37,9 +39,8 @@ def _gather_height_sizes(tensor: torch.Tensor) -> list[int]: if world_size == 1: return [tensor.shape[-2]] local_height = torch.tensor([tensor.shape[-2]], dtype=torch.int64, device=tensor.device) - gathered = [torch.empty_like(local_height) for _ in range(world_size)] - dist.all_gather(gathered, local_height) - return [int(height.item()) for height in gathered] + gathered = all_gather_stacked(local_height, world_size=world_size) + return gathered.flatten().tolist() def _gather_height(tensor: torch.Tensor) -> torch.Tensor: @@ -47,15 +48,19 @@ def _gather_height(tensor: torch.Tensor) -> torch.Tensor: if world_size == 1: return tensor heights = _gather_height_sizes(tensor) + memory_format = _height_memory_format(tensor) max_height = max(heights) if tensor.shape[-2] < max_height: shape = list(tensor.shape) shape[-2] = max_height - tensor.shape[-2] tensor = torch.cat([tensor, tensor.new_zeros(shape)], dim=-2) - tensor = tensor.contiguous() - gathered = [torch.empty_like(tensor) for _ in range(world_size)] - dist.all_gather(gathered, tensor) - return torch.cat([shard[..., :height, :] for shard, height in zip(gathered, heights)], dim=-2) + gather_input = tensor.movedim(-2, 0).contiguous() + gathered = all_gather_stacked(gather_input, world_size=world_size) + if all(height == max_height for height in heights): + merged = gathered.flatten(0, 1) + else: + merged = torch.cat([shard[:height] for shard, height in zip(gathered, heights)], dim=0) + return merged.movedim(0, -2).contiguous(memory_format=memory_format) def _ensure_halo_buffer(buffer: torch.Tensor | None, reference: torch.Tensor) -> torch.Tensor: @@ -78,21 +83,25 @@ def _exchange_height_halo(module: nn.Module, tensor: torch.Tensor, halo_size: in rank = _spatial_rank() top = tensor[..., :halo_size, :] bottom = tensor[..., -halo_size:, :] + module._halo_send_top = _ensure_halo_buffer(module._halo_send_top, top) + module._halo_send_bottom = _ensure_halo_buffer(module._halo_send_bottom, bottom) module._halo_recv_top = _ensure_halo_buffer(module._halo_recv_top, top) module._halo_recv_bottom = _ensure_halo_buffer(module._halo_recv_bottom, bottom) + module._halo_send_top.copy_(top) + module._halo_send_bottom.copy_(bottom) operations = [] if rank > 0: operations.extend( [ dist.P2POp(dist.irecv, module._halo_recv_top, rank - 1), - dist.P2POp(dist.isend, top.contiguous(memory_format=_height_memory_format(top)), rank - 1), + dist.P2POp(dist.isend, module._halo_send_top, rank - 1), ] ) if rank < world_size - 1: operations.extend( [ - dist.P2POp(dist.isend, bottom.contiguous(memory_format=_height_memory_format(bottom)), rank + 1), + dist.P2POp(dist.isend, module._halo_send_bottom, rank + 1), dist.P2POp(dist.irecv, module._halo_recv_bottom, rank + 1), ] ) @@ -158,6 +167,8 @@ def __init__(self, source: nn.Conv2d) -> None: if source.padding[0] != self._height_halo_size: raise ValueError("VAE spatial Conv2d requires symmetric height padding") self._width_padding = source.padding[1] + self._halo_send_top: torch.Tensor | None = None + self._halo_send_bottom: torch.Tensor | None = None self._halo_recv_top: torch.Tensor | None = None self._halo_recv_bottom: torch.Tensor | None = None self.train(source.training) diff --git a/telefuser/models/lingbot_video_dit.py b/telefuser/models/lingbot_video_dit.py index 9af27de..510f368 100644 --- a/telefuser/models/lingbot_video_dit.py +++ b/telefuser/models/lingbot_video_dit.py @@ -15,7 +15,8 @@ from torch import nn from telefuser.core.model_registry import register_model_config -from telefuser.distributed.ulysses_comm import ulysses_all_to_all_split_cat +from telefuser.distributed.collectives import all_gather_cat +from telefuser.distributed.ulysses_comm import ulysses_gather_heads, ulysses_scatter_heads from telefuser.ops.attention import attention from telefuser.utils.model_weight import hash_state_dict_keys @@ -154,26 +155,12 @@ def forward( group is not None and dist.is_available() and dist.is_initialized() and dist.get_world_size(group) > 1 ) if use_ulysses: - world_size = dist.get_world_size(group) - local_heads = self.num_heads // world_size - query = ulysses_all_to_all_split_cat( - query.reshape(batch, sequence, self.num_heads * self.head_dim), - group, - scatter_dim=2, - gather_dim=1, - ).view(batch, sequence * world_size, local_heads, self.head_dim) - key = ulysses_all_to_all_split_cat( - key.reshape(batch, sequence, self.num_heads * self.head_dim), - group, - scatter_dim=2, - gather_dim=1, - ).view(batch, sequence * world_size, local_heads, self.head_dim) - value = ulysses_all_to_all_split_cat( - value.reshape(batch, sequence, self.num_heads * self.head_dim), - group, - scatter_dim=2, - gather_dim=1, - ).view(batch, sequence * world_size, local_heads, self.head_dim) + query_wait = ulysses_scatter_heads(query, group) + key_wait = ulysses_scatter_heads(key, group) + value_wait = ulysses_scatter_heads(value, group) + query = query_wait() + key = key_wait() + value = value_wait() output = attention( query.transpose(1, 2), key.transpose(1, 2), @@ -188,12 +175,7 @@ def forward( raise RuntimeError("LingBot attention does not support log-sum-exp outputs") output = output.transpose(1, 2) if use_ulysses: - output = ulysses_all_to_all_split_cat( - output.reshape(batch, sequence * world_size, local_heads * self.head_dim), - group, - scatter_dim=1, - gather_dim=2, - ).view(batch, sequence, self.num_heads, self.head_dim) + output = ulysses_gather_heads(output, group, num_heads=self.num_heads)() return self.to_out(output.reshape(batch, sequence, -1).to(hidden_states.dtype)) @@ -419,11 +401,7 @@ def _ulysses_shard_joint( def _ulysses_gather_sequence(self, local: torch.Tensor) -> torch.Tensor: """Gather equal local token slices in rank order without changing their layout.""" - if self._ulysses_world_size() == 1: - return local - gathered = [torch.empty_like(local) for _ in range(self._ulysses_world_size())] - dist.all_gather(gathered, local.contiguous(), group=self.ulysses_group) - return torch.cat(gathered, dim=1) + return all_gather_cat(local, dim=1, group=self.ulysses_group, world_size=self._ulysses_world_size()) def forward( self, diff --git a/telefuser/models/liveact_dit.py b/telefuser/models/liveact_dit.py index 92d05d2..c5e197d 100644 --- a/telefuser/models/liveact_dit.py +++ b/telefuser/models/liveact_dit.py @@ -376,9 +376,12 @@ def forward_sp( v = self.v(x).view(b, s, n, d) # Ulysses scatter heads: [B, S/N, H, D] -> [B, S, H/N, D] - q = ulysses_scatter_heads(q, self.ulysses_group)() - k = ulysses_scatter_heads(k, self.ulysses_group)() - v = ulysses_scatter_heads(v, self.ulysses_group)() + q_wait = ulysses_scatter_heads(q, self.ulysses_group) + k_wait = ulysses_scatter_heads(k, self.ulysses_group) + v_wait = ulysses_scatter_heads(v, self.ulysses_group) + q = q_wait() + k = k_wait() + v = v_wait() k_cache, v_cache = kv_cache.load(x.device, torch.bfloat16) diff --git a/telefuser/models/video_projector.py b/telefuser/models/video_projector.py index d8d842e..6d6629f 100644 --- a/telefuser/models/video_projector.py +++ b/telefuser/models/video_projector.py @@ -7,6 +7,8 @@ from einops import rearrange, repeat from tqdm import tqdm +from telefuser.distributed.collectives import all_reduce_sum_ + CACHE_T = 2 @@ -386,8 +388,7 @@ def tile_stream_forward( weight = weights[layer_idx] all_value = all_values[layer_idx] if self.parallelism > 1: - dist.all_reduce(all_value) - dist.all_reduce(weight) + all_reduce_sum_((all_value, weight)) weight[weight == 0] = 1 averaged = all_value / weight.unsqueeze(-1) averaged = averaged.view(1, -1, self.linear_layers[0].out_features).cpu() diff --git a/telefuser/models/wan22_video_vae.py b/telefuser/models/wan22_video_vae.py index a199487..9cfbf44 100644 --- a/telefuser/models/wan22_video_vae.py +++ b/telefuser/models/wan22_video_vae.py @@ -6,6 +6,7 @@ from einops import rearrange from telefuser.core.base_model import BaseModel +from telefuser.distributed.collectives import all_gather_stacked, all_reduce_sum_ # Import shared components from existing VAE from .wan_video_vae import ( @@ -951,8 +952,6 @@ def encode_dist_2d( device: torch.device, ) -> torch.Tensor: """Encode video with true 2D spatial splitting.""" - import torch.distributed as dist - spatial_ratio = self.upsampling_factor # 16 for Wan2.2 padding_latent = 1 @@ -980,9 +979,7 @@ def encode_dist_2d( encoded_chunk, cur_rank_h, cur_rank_w, world_size_h, world_size_w, chunk_h, chunk_w ) - world_size_total = world_size_h * world_size_w - full_encoded = [torch.empty_like(encoded_chunk) for _ in range(world_size_total)] - dist.all_gather(full_encoded, encoded_chunk) + full_encoded = list(all_gather_stacked(encoded_chunk, world_size=world_size_h * world_size_w).unbind(0)) encoded = self._reconstruct_2d(full_encoded, world_size_h, world_size_w, dim=3) return encoded.squeeze(0).cpu() @@ -997,8 +994,6 @@ def decode_dist_2d( device: torch.device, ) -> torch.Tensor: """Decode latent with true 2D spatial splitting.""" - import torch.distributed as dist - spatial_ratio = self.upsampling_factor # 16 padding_latent = 2 @@ -1022,9 +1017,7 @@ def decode_dist_2d( decoded_chunk, cur_rank_h, cur_rank_w, world_size_h, world_size_w, chunk_h_output, chunk_w_output ) - world_size_total = world_size_h * world_size_w - full_decoded = [torch.empty_like(decoded_chunk) for _ in range(world_size_total)] - dist.all_gather(full_decoded, decoded_chunk) + full_decoded = list(all_gather_stacked(decoded_chunk, world_size=world_size_h * world_size_w).unbind(0)) decoded = self._reconstruct_2d(full_decoded, world_size_h, world_size_w, dim=3) return decoded.squeeze(0).cpu().clamp_(-1, 1) @@ -1225,8 +1218,7 @@ def tiled_encode( weight[:, :, :, target_h:target_h_end, target_w:target_w_end] += mask if self.parallelism > 1 and dist.is_initialized(): - dist.all_reduce(values) - dist.all_reduce(weight) + all_reduce_sum_((values, weight)) values = values / weight values = values.float() return values @@ -1297,8 +1289,7 @@ def tiled_decode( weight[:, :, :, target_h:target_h_end, target_w:target_w_end] += mask if self.parallelism > 1 and dist.is_initialized(): - dist.all_reduce(values) - dist.all_reduce(weight) + all_reduce_sum_((values, weight)) values = values / weight values = values.cpu() # unpatchify is already called in VideoVAE.decode(), output is already RGB diff --git a/telefuser/models/wan_video_vae.py b/telefuser/models/wan_video_vae.py index 6ffa7f2..22b4bc7 100644 --- a/telefuser/models/wan_video_vae.py +++ b/telefuser/models/wan_video_vae.py @@ -10,6 +10,7 @@ from tqdm import tqdm from telefuser.core.base_model import BaseModel +from telefuser.distributed.collectives import all_gather_stacked, all_reduce_sum_ from telefuser.distributed.vae_spatial import ( _SpatialParallelConv2d, _gather_height, @@ -157,6 +158,8 @@ def __init__(self, source: CausalConv3d) -> None: 2 * temporal_padding, 0, ) + self._halo_send_top: torch.Tensor | None = None + self._halo_send_bottom: torch.Tensor | None = None self._halo_recv_top: torch.Tensor | None = None self._halo_recv_bottom: torch.Tensor | None = None self.train(source.training) @@ -506,12 +509,20 @@ def __init__( CausalConv3d(out_dim, 3, 3, padding=1), ) - def forward(self, x: torch.Tensor, feat_cache: list | None = None, feat_idx: list | None = None) -> torch.Tensor: + def forward( + self, + x: torch.Tensor, + feat_cache: list | None = None, + feat_idx: list | None = None, + input_global_height: int | None = None, + ) -> torch.Tensor: """Forward pass with list-based feature caching.""" expected_height = None if self._spatial_parallel: - expected_height = x.shape[-2] * (2**self._spatial_upsample_count) - x = _split_height(x) + global_height = x.shape[-2] if input_global_height is None else input_global_height + expected_height = global_height * (2**self._spatial_upsample_count) + if input_global_height is None: + x = _split_height(x) if feat_cache is not None and feat_idx is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() @@ -1064,9 +1075,7 @@ def encode_dist_2d( ) # Gather all chunks - world_size_total = world_size_h * world_size_w - full_encoded = [torch.empty_like(encoded_chunk) for _ in range(world_size_total)] - dist.all_gather(full_encoded, encoded_chunk) + full_encoded = list(all_gather_stacked(encoded_chunk, world_size=world_size_h * world_size_w).unbind(0)) # Reconstruct full latent encoded = self._reconstruct_2d(full_encoded, world_size_h, world_size_w, dim=3) @@ -1128,9 +1137,7 @@ def decode_dist_2d( ) # Gather all chunks - world_size_total = world_size_h * world_size_w - full_decoded = [torch.empty_like(decoded_chunk) for _ in range(world_size_total)] - dist.all_gather(full_decoded, decoded_chunk) + full_decoded = list(all_gather_stacked(decoded_chunk, world_size=world_size_h * world_size_w).unbind(0)) # Reconstruct full video decoded = self._reconstruct_2d(full_decoded, world_size_h, world_size_w, dim=3) @@ -1326,8 +1333,7 @@ def tiled_decode( weight[:, :, :, target_h:target_h_end, target_w:target_w_end] += mask if self.parallelism > 1: - dist.all_reduce(values) - dist.all_reduce(weight) + all_reduce_sum_((values, weight)) values = values / weight # Move to CPU to reduce VRAM usage (video output is large) values = values.cpu().clamp_(-1, 1) @@ -1390,8 +1396,7 @@ def tiled_encode( weight[:, :, :, target_h:target_h_end, target_w:target_w_end] += mask if self.parallelism > 1: - dist.all_reduce(values) - dist.all_reduce(weight) + all_reduce_sum_((values, weight)) values = values / weight return values @@ -1564,6 +1569,7 @@ def cached_decode_withflag( is_first_clip: bool, is_last_clip: bool, decode_state: WanVideoVAEStreamingDecodeState | None = None, + input_global_height: int | None = None, ) -> torch.Tensor: """Decode with persistent feature cache for streaming generation. @@ -1613,6 +1619,7 @@ def cached_decode_withflag( # Decode frame-by-frame with cache iter_ = z.shape[2] x = self.model.conv2(z) + decoder_kwargs = {"input_global_height": input_global_height} if input_global_height is not None else {} for i in range(iter_): feat_idx[0] = 0 # Reset index for each frame @@ -1621,12 +1628,14 @@ def cached_decode_withflag( x[:, :, i : i + 1, :, :], feat_cache=feat_cache, feat_idx=feat_idx, + **decoder_kwargs, ) else: out_ = self.model.decoder( x[:, :, i : i + 1, :, :], feat_cache=feat_cache, feat_idx=feat_idx, + **decoder_kwargs, ) out = torch.cat([out, out_], 2) diff --git a/telefuser/pipelines/lingbot_world_fast/pipeline.py b/telefuser/pipelines/lingbot_world_fast/pipeline.py index cba5275..70a68ae 100644 --- a/telefuser/pipelines/lingbot_world_fast/pipeline.py +++ b/telefuser/pipelines/lingbot_world_fast/pipeline.py @@ -375,6 +375,7 @@ def init(self, module_manager: ModuleManager, config: LingBotWorldFastPipelineCo dit_runtime_config.parallel_config.timeout, vae_decode_config.parallel_config.timeout, ), + shard_dim=-2 if vae_decode_config.parallel_config.world_size > 1 else None, ) self._worker_tensor_channels.append(latent_channel) self.denoise_stage = ( @@ -662,12 +663,15 @@ def close(self) -> None: self._streaming_runtime = None if streaming_runtime is not None: streaming_runtime.close() + vae_decode_worker = getattr(self, "vae_decode_worker", None) + if isinstance(vae_decode_worker, ParallelWorker): + vae_decode_worker.close() denoise_stage = getattr(self, "denoise_stage", None) if isinstance(denoise_stage, ParallelWorker): denoise_stage.close() - for vae_worker in (getattr(self, "vae_encode_worker", None), getattr(self, "vae_decode_worker", None)): - if isinstance(vae_worker, ParallelWorker): - vae_worker.close() + vae_encode_worker = getattr(self, "vae_encode_worker", None) + if isinstance(vae_encode_worker, ParallelWorker): + vae_encode_worker.close() for channel in getattr(self, "_worker_tensor_channels", ()): channel.close() diff --git a/telefuser/pipelines/lingbot_world_fast/streaming.py b/telefuser/pipelines/lingbot_world_fast/streaming.py index 7c1d828..46938f1 100644 --- a/telefuser/pipelines/lingbot_world_fast/streaming.py +++ b/telefuser/pipelines/lingbot_world_fast/streaming.py @@ -4,6 +4,7 @@ import threading import time +from collections import deque from collections.abc import Callable from contextlib import nullcontext from dataclasses import dataclass, field @@ -28,6 +29,7 @@ ) from telefuser.utils.logging import logger from telefuser.worker.parallel_worker import ParallelWorker +from telefuser.worker.tensor_channel import WorkerTensorRef from .session import LingBotWorldFastGenerationSession, LingBotWorldFastSessionStatus @@ -57,6 +59,17 @@ class _LingBotStreamingSessionEntry: chunk_profiles: dict[int, dict[str, object]] = field(default_factory=dict) +@dataclass +class _DirectTensorTransfer: + """One channel transfer awaiting consumption or cancellation cleanup.""" + + session_id: str + key: tuple[str, int] + value: object + tensor_count: int + cancelled: bool = False + + class LingBotWorldFastStreamingRuntime: """Own the single actor graph shared by all sessions of one pipeline.""" @@ -64,6 +77,8 @@ def __init__(self, pipeline: LingBotWorldFastPipeline) -> None: self.pipeline = pipeline self._lock = threading.RLock() self._sessions: dict[str, _LingBotStreamingSessionEntry] = {} + self._direct_condition_transfers: deque[_DirectTensorTransfer] = deque() + self._direct_latent_transfers: deque[_DirectTensorTransfer] = deque() self._closed = False self._serialize_dit_decode = self._dit_decode_devices_overlap() self._dit_decode_lock = threading.Lock() @@ -365,6 +380,146 @@ def _dit_decode_devices_overlap(self) -> bool: decode_devices = self._runtime_device_ids(config.vae_decode_config) return not dit_devices.isdisjoint(decode_devices) + @staticmethod + def _direct_transfer_refs(value: object) -> tuple[WorkerTensorRef, ...]: + refs: list[WorkerTensorRef] = [] + seen: set[WorkerTensorRef] = set() + + def visit(item: object) -> None: + if isinstance(item, WorkerTensorRef): + if item not in seen: + seen.add(item) + refs.append(item) + elif isinstance(item, dict): + for child in item.values(): + visit(child) + elif isinstance(item, tuple | list): + for child in item: + visit(child) + + visit(value) + return tuple(refs) + + @staticmethod + def _direct_transfer_key(refs: tuple[WorkerTensorRef, ...]) -> tuple[str, int]: + keys = {(ref.channel_id, ref.transfer_id) for ref in refs} + if len(keys) != 1: + raise RuntimeError(f"Direct tensor artifact must contain one channel transfer, got {sorted(keys)}") + return next(iter(keys)) + + def _track_direct_transfer( + self, + transfers: deque[_DirectTensorTransfer], + session_id: str, + value: object, + ) -> None: + refs = self._direct_transfer_refs(value) + if not refs: + return + transfer = _DirectTensorTransfer( + session_id=session_id, + key=self._direct_transfer_key(refs), + value=value, + tensor_count=len(refs), + ) + with self._lock: + transfers.append(transfer) + + def _consume_direct_transfer( + self, + transfers: deque[_DirectTensorTransfer], + session_id: str, + value: object, + worker: object, + label: str, + ) -> None: + """Retire one consumed transfer and older entries skipped by channel receive.""" + refs = self._direct_transfer_refs(value) + if not refs: + return + key = self._direct_transfer_key(refs) + with self._lock: + match_index = None + for index, transfer in enumerate(transfers): + if transfer.key == key: + if transfer.session_id != session_id: + raise RuntimeError( + f"Direct {label} transfer {key[1]} belongs to {transfer.session_id!r}, not {session_id!r}" + ) + match_index = index + break + if not transfer.cancelled: + raise RuntimeError(f"Direct {label} transfer {key[1]} overtook live transfer {transfer.key[1]}") + if match_index is None: + raise RuntimeError(f"Unknown direct {label} transfer {key[1]}") + for _ in range(match_index + 1): + transfers.popleft() + self._discard_cancelled_direct_transfers(transfers, worker, label) + + def _cancel_direct_transfers( + self, + transfers: deque[_DirectTensorTransfer], + session_id: str, + worker: object, + label: str, + ) -> None: + with self._lock: + for transfer in transfers: + if transfer.session_id == session_id: + transfer.cancelled = True + self._discard_cancelled_direct_transfers(transfers, worker, label) + + def _discard_direct_transfer( + self, + transfers: deque[_DirectTensorTransfer], + session_id: str, + value: object, + worker: object, + label: str, + ) -> None: + refs = self._direct_transfer_refs(value) + if not refs: + return + key = self._direct_transfer_key(refs) + with self._lock: + for transfer in transfers: + if transfer.key != key: + continue + if transfer.session_id != session_id: + raise RuntimeError( + f"Direct {label} transfer {key[1]} belongs to {transfer.session_id!r}, not {session_id!r}" + ) + transfer.cancelled = True + break + else: + raise RuntimeError(f"Unknown direct {label} transfer {key[1]}") + self._discard_cancelled_direct_transfers(transfers, worker, label) + + def _discard_cancelled_direct_transfers( + self, + transfers: deque[_DirectTensorTransfer], + worker: object, + label: str, + ) -> None: + """Drain only the cancelled FIFO prefix from every consumer rank.""" + while True: + with self._lock: + if not transfers or not transfers[0].cancelled: + return + transfer = transfers[0] + if not isinstance(worker, ParallelWorker): + raise RuntimeError(f"Direct LingBot {label} cleanup requires a ParallelWorker consumer") + discarded = worker.discard_tensor_refs(transfer.value, sync=True) + if discarded != transfer.tensor_count: + raise RuntimeError( + f"Direct {label} cleanup discarded {discarded} of {transfer.tensor_count} tensors " + f"for transfer {transfer.key[1]}" + ) + with self._lock: + if not transfers or transfers[0] is not transfer: + raise RuntimeError(f"Direct {label} transfer order changed during cancellation cleanup") + transfers.popleft() + def _entry_for_context(self, context: StreamingSessionContext) -> _LingBotStreamingSessionEntry: with self._lock: try: @@ -396,6 +551,12 @@ def _release_decode_session( ) -> None: del reason entry = self._entry_for_context(context) + self._cancel_direct_transfers( + self._direct_latent_transfers, + context.session_id, + self.pipeline.vae_decode_worker, + "latent", + ) cache_handle = entry.runtime.cache_handle if cache_handle is None: return @@ -410,6 +571,12 @@ def _release_denoise_session( ) -> None: del reason entry = self._entry_for_context(context) + self._cancel_direct_transfers( + self._direct_condition_transfers, + context.session_id, + self.pipeline.denoise_stage, + "condition", + ) cache_handle = entry.runtime.cache_handle if cache_handle is None: return @@ -456,6 +623,7 @@ def _encode_inputs(self, invocation: StreamingStageInvocation) -> tuple[tuple[ob def _encode_outputs(self, value: dict[str, object], invocation: StreamingStageInvocation) -> dict[str, object]: entry = self._entry_for_invocation(invocation) index = invocation.key.sequence_id + self._track_direct_transfer(self._direct_condition_transfers, invocation.key.session_id, value) self.pipeline._notify_progress(entry.progress_callback, "condition_chunk_encoded", index=index) if index == 0: entry.runtime.condition_image = None @@ -486,8 +654,16 @@ def _denoise(self, invocation: StreamingStageInvocation) -> dict[str, object]: entry = self._entry_for_invocation(invocation) runtime = entry.runtime index = invocation.key.sequence_id + condition = invocation.inputs["condition"] cached_latent = runtime.world_kv_cached_latents.pop(index, None) if runtime.world_kv_cached_latents else None if cached_latent is not None: + self._discard_direct_transfer( + self._direct_condition_transfers, + invocation.key.session_id, + condition, + self.pipeline.denoise_stage, + "condition", + ) self.pipeline._notify_progress(entry.progress_callback, "world_kv_cache_hit", index=index) advance = self.pipeline.denoise_stage.advance_noise(cache_handle=runtime.cache_handle) if callable(advance): @@ -498,24 +674,34 @@ def _denoise(self, invocation: StreamingStageInvocation) -> dict[str, object]: kwargs = self._denoise_kwargs(invocation) lock = self._dit_decode_lock if self._serialize_dit_decode else nullcontext() lock_started_at = time.perf_counter() - with lock: - worker_started_at = time.perf_counter() - result = self.pipeline.denoise_stage.denoise_and_update_cache(**kwargs) - submit_finished_at = time.perf_counter() - if callable(result): - result = result() - worker_finished_at = time.perf_counter() - if isinstance(result, tuple): - latent, profile = result - profile["denoise_lock_wait_seconds"] = worker_started_at - lock_started_at - profile["denoise_submit_seconds"] = submit_finished_at - worker_started_at - profile["denoise_result_wait_seconds"] = worker_finished_at - submit_finished_at - profile["denoise_worker_seconds"] = worker_finished_at - worker_started_at - with self._lock: - entry.chunk_profiles.setdefault(index, {}).update(profile) - else: - latent = result + try: + with lock: + worker_started_at = time.perf_counter() + result = self.pipeline.denoise_stage.denoise_and_update_cache(**kwargs) + submit_finished_at = time.perf_counter() + if callable(result): + result = result() + worker_finished_at = time.perf_counter() + finally: + self._consume_direct_transfer( + self._direct_condition_transfers, + invocation.key.session_id, + condition, + self.pipeline.denoise_stage, + "condition", + ) + if isinstance(result, tuple): + latent, profile = result + profile["denoise_lock_wait_seconds"] = worker_started_at - lock_started_at + profile["denoise_submit_seconds"] = submit_finished_at - worker_started_at + profile["denoise_result_wait_seconds"] = worker_finished_at - submit_finished_at + profile["denoise_worker_seconds"] = worker_finished_at - worker_started_at + with self._lock: + entry.chunk_profiles.setdefault(index, {}).update(profile) + else: + latent = result self.pipeline._notify_progress(entry.progress_callback, "chunk_denoised", index=index) + self._track_direct_transfer(self._direct_latent_transfers, invocation.key.session_id, latent) if runtime.world_kv_binding is not None: try: runtime.world_kv_binding.on_chunk_finalized(runtime, index, latent) @@ -534,13 +720,17 @@ def _decode_inputs(self, invocation: StreamingStageInvocation) -> tuple[tuple[ob index=index, device=str(self.pipeline.vae_device), ) - return (), { + latent = invocation.inputs["latent"] + kwargs = { "cache_handle": runtime.cache_handle, - "latents": invocation.inputs["latent"], + "latents": latent, "is_first_clip": index == 0, "is_last_clip": index == runtime.chunk_count - 1, "_benchmark_profile": runtime.config.benchmark_metrics, } + if isinstance(latent, WorkerTensorRef) and latent.shard_dim == len(latent.shape) - 2: + kwargs["_global_latent_height"] = latent.shape[-2] + return (), kwargs def _decode_outputs(self, value: torch.Tensor, invocation: StreamingStageInvocation) -> dict[str, object]: entry = self._entry_for_invocation(invocation) @@ -560,15 +750,25 @@ def _decode_outputs(self, value: torch.Tensor, invocation: StreamingStageInvocat @torch.inference_mode() def _decode(self, invocation: StreamingStageInvocation) -> dict[str, object]: args, kwargs = self._decode_inputs(invocation) + latent = kwargs["latents"] lock = self._dit_decode_lock if self._serialize_dit_decode else nullcontext() lock_started_at = time.perf_counter() - with lock: - worker_started_at = time.perf_counter() - result = self.pipeline.vae_decode_worker.decode_chunk(*args, **kwargs) - submit_finished_at = time.perf_counter() - if callable(result): - result = result() - worker_finished_at = time.perf_counter() + try: + with lock: + worker_started_at = time.perf_counter() + result = self.pipeline.vae_decode_worker.decode_chunk(*args, **kwargs) + submit_finished_at = time.perf_counter() + if callable(result): + result = result() + worker_finished_at = time.perf_counter() + finally: + self._consume_direct_transfer( + self._direct_latent_transfers, + invocation.key.session_id, + latent, + self.pipeline.vae_decode_worker, + "latent", + ) if isinstance(result, tuple): value, profile = result profile["decode_lock_wait_seconds"] = worker_started_at - lock_started_at diff --git a/telefuser/pipelines/lingbot_world_fast/vae_stage.py b/telefuser/pipelines/lingbot_world_fast/vae_stage.py index 270eb58..28fb496 100644 --- a/telefuser/pipelines/lingbot_world_fast/vae_stage.py +++ b/telefuser/pipelines/lingbot_world_fast/vae_stage.py @@ -379,6 +379,7 @@ def decode_chunk( is_first_clip: bool, is_last_clip: bool, _benchmark_profile: bool = False, + _global_latent_height: int | None = None, ) -> torch.Tensor | None | tuple[torch.Tensor | None, dict[str, float]]: """Decode one latent chunk and return CPU frame tensors.""" state = self._cache_registry[cache_handle] @@ -393,6 +394,7 @@ def decode_chunk( is_first_clip=is_first_clip, is_last_clip=is_last_clip, decode_state=state.decoder_state, + input_global_height=_global_latent_height, ) profile = None if decode_end is not None: diff --git a/telefuser/worker/parallel_worker.py b/telefuser/worker/parallel_worker.py index 394858b..b814e0b 100644 --- a/telefuser/worker/parallel_worker.py +++ b/telefuser/worker/parallel_worker.py @@ -32,6 +32,9 @@ from telefuser.metrics import StageMetricContext +_DISCARD_TENSOR_REFS = "__telefuser_discard_tensor_refs__" + + def to_device(data: Any, device: str | torch.device) -> Any: """Recursively move data to target device.""" if isinstance(data, dict): @@ -121,6 +124,16 @@ def _worker_loop( if name == "exit": logger.info(f"parallel worker {stage.name} on rank {rank} exits") break + if name == _DISCARD_TENSOR_REFS: + discarded = 0 + for channel in tensor_input_channels: + if channel.contains_ref(args): + discarded += channel.discard(args, rank=rank) + if world_size > 1: + dist.barrier() + if world_size == 1 or rank == 0: + queue_out.put(discarded) + continue if not hasattr(stage, name): raise AttributeError(f'{stage.__class__.__name__} has no attribute "{name}"') stage_inputs = (args, kwargs) @@ -157,6 +170,8 @@ def _worker_loop( args = None kwargs = None current_platform.synchronize() + for channel in tensor_input_channels: + channel.release_local_cuda_ipc() gc.collect() current_platform.empty_cache() current_platform.ipc_collect() @@ -332,6 +347,16 @@ def put_data(self, data: Any) -> None: for q in self.queue_in: q.put(data) + def discard_tensor_refs(self, value: Any, *, sync: bool = False) -> int | Callable[[], int]: + """Release direct-channel tensors that will not be passed to the stage.""" + self._ensure_usable() + self.put_data([_DISCARD_TENSOR_REFS, (value,), {}, False]) + + def wait() -> int: + return int(self._wait_result(_DISCARD_TENSOR_REFS)) + + return wait() if sync else wait + def __call__(self, *args: Any, **kwargs: Any) -> Any | Callable[[], Any]: """Submit __call__ task to all workers.""" self._ensure_usable() diff --git a/telefuser/worker/ray_worker.py b/telefuser/worker/ray_worker.py index 9bb63ef..5b7cb7f 100644 --- a/telefuser/worker/ray_worker.py +++ b/telefuser/worker/ray_worker.py @@ -6,7 +6,6 @@ from __future__ import annotations -import os from typing import Any import torch @@ -46,11 +45,17 @@ def _setup_resources(self) -> None: logger.info(f"has device {current_platform.device_type}") if gpu_config.num_gpus > 0: - gpu_ids = list(range(gpu_config.num_gpus)) - os.environ[current_platform.device_control_env_var] = ",".join(map(str, gpu_ids)) - logger.info(f"RayWorker {self.worker_id} use device: {gpu_ids}") + visible_gpu_count = current_platform.device_count() + if visible_gpu_count < gpu_config.num_gpus: + raise RuntimeError( + f"Ray assigned {visible_gpu_count} visible GPUs to {self.worker_id}, " + f"but the stage requires {gpu_config.num_gpus}" + ) + logical_gpu_ids = list(range(gpu_config.num_gpus)) + current_platform.set_device(logical_gpu_ids[0]) + logger.info(f"RayWorker {self.worker_id} use Ray-assigned logical devices: {logical_gpu_ids}") if gpu_config.memory_limit > 0 and current_platform.device_type == "cuda": - torch.cuda.set_per_process_memory_fraction(gpu_config.memory_limit) + torch.cuda.set_per_process_memory_fraction(gpu_config.memory_limit, device=logical_gpu_ids[0]) logger.info(f"RayWorker {self.worker_id} GPU memory limit: {gpu_config.memory_limit}") if self.ray_config.memory_gb > 0: diff --git a/telefuser/worker/tensor_channel.py b/telefuser/worker/tensor_channel.py index 6046a2a..22f9011 100644 --- a/telefuser/worker/tensor_channel.py +++ b/telefuser/worker/tensor_channel.py @@ -2,14 +2,18 @@ from __future__ import annotations +import time import uuid -from dataclasses import dataclass +from dataclasses import dataclass, field +from dataclasses import replace as dataclass_replace from multiprocessing.queues import SimpleQueue from typing import Any import torch import torch.multiprocessing as mp +_MAX_CUDA_IPC_POOL_PROFILES = 8 + @dataclass(frozen=True) class WorkerTensorRef: @@ -22,26 +26,94 @@ class WorkerTensorRef: dtype: str source_device: str nbytes: int + shard_dim: int | None = None + _cuda_payloads: tuple[_CudaIpcPayload, ...] | None = field( + default=None, + compare=False, + hash=False, + repr=False, + ) + + +@dataclass(frozen=True) +class _CudaIpcPayload: + """Plain metadata for one view inside a producer-owned persistent CUDA pool.""" + + ref: WorkerTensorRef + profile_index: int + slot_id: int + handle: tuple[Any, ...] + ready_event_handle: bytes + byte_offset: int + shape: tuple[int, ...] + stride: tuple[int, ...] + dtype: torch.dtype + + +@dataclass +class _CudaIpcSlot: + tensor: torch.Tensor + pending_ranks: set[int] + completion_ranks: set[int] + ready_event: torch.cuda.Event + ready_event_handle: bytes + was_used: bool = False + transfer_key: tuple[int, int] | None = None + + +@dataclass +class _CudaIpcPool: + profile_index: int + storage: torch.Tensor + handle: tuple[Any, ...] + slots: list[_CudaIpcSlot] + next_slot_id: int = 0 class WorkerTensorChannel: """Point-to-point tensor path that bypasses the parent process. - The producer places tensors directly onto one queue per consumer rank. CUDA - tensors travel as CUDA IPC handles; CPU tensors use multiprocessing shared - memory. The parent process receives only :class:`WorkerTensorRef` objects. + CUDA tensors use persistent IPC pools whose rank-local metadata travels in + :class:`WorkerTensorRef`. CPU tensors and dynamic-profile CUDA fallbacks use + one multiprocessing queue per consumer rank. The parent never materializes + tensor contents. """ - def __init__(self, consumer_world_size: int, *, timeout: int = 600) -> None: + def __init__( + self, + consumer_world_size: int, + *, + timeout: int = 600, + shard_dim: int | None = None, + cuda_ipc_slots: int = 2, + ) -> None: if consumer_world_size < 1: raise ValueError("consumer_world_size must be at least one") if timeout < 1: raise ValueError("timeout must be at least one second") + if cuda_ipc_slots < 1: + raise ValueError("cuda_ipc_slots must be at least one") spawn_ctx = mp.get_context("spawn") self.channel_id = uuid.uuid4().hex self.consumer_world_size = consumer_world_size self.timeout = timeout + self.shard_dim = shard_dim + self.cuda_ipc_slots = cuda_ipc_slots self._queues: tuple[SimpleQueue, ...] = tuple(spawn_ctx.SimpleQueue() for _ in range(consumer_world_size)) + self._ack_generations = spawn_ctx.Array( + "q", + _MAX_CUDA_IPC_POOL_PROFILES * cuda_ipc_slots * consumer_world_size, + lock=False, + ) + self._completion_queues: tuple[SimpleQueue, ...] = tuple( + spawn_ctx.SimpleQueue() for _ in range(consumer_world_size) + ) + self._cuda_pools: dict[tuple[int, tuple[int, ...], torch.dtype, str], _CudaIpcPool] = {} + self._cuda_storage_cache: dict[tuple[bytes, int], torch.UntypedStorage] = {} + self._cuda_event_cache: dict[tuple[bytes, int], torch.cuda.Event] = {} + self._cuda_consumer_completion_events: dict[tuple[int, int], torch.cuda.Event] = {} + self._cuda_completion_handles: dict[tuple[int, int, int], bytes] = {} + self._cuda_producer_completion_events: dict[tuple[int, int, int, int], torch.cuda.Event] = {} self._next_transfer_id = 0 self._producer_bound = False self._consumer_bound = False @@ -80,6 +152,7 @@ def replace(item: Any) -> Any: existing = sent.get(id(item)) if existing is not None: return existing + shard_dim = self._normalize_shard_dim(item) ref = WorkerTensorRef( channel_id=self.channel_id, transfer_id=transfer_id, @@ -88,11 +161,30 @@ def replace(item: Any) -> Any: dtype=str(item.dtype), source_device=str(item.device), nbytes=item.numel() * item.element_size(), + shard_dim=shard_dim, ) tensor_index += 1 + if item.device.type == "cuda": + rank_payloads = self._stage_cuda_tensor(ref, item) + if rank_payloads is not None: + ref = dataclass_replace(ref, _cuda_payloads=rank_payloads) + else: + rank_tensors = ( + (item,) * self.consumer_world_size + if shard_dim is None + else torch.tensor_split(item, self.consumer_world_size, dim=shard_dim) + ) + for queue, rank_tensor in zip(self._queues, rank_tensors): + queue.put((ref, rank_tensor)) + else: + rank_tensors = ( + (item,) * self.consumer_world_size + if shard_dim is None + else torch.tensor_split(item, self.consumer_world_size, dim=shard_dim) + ) + for queue, rank_tensor in zip(self._queues, rank_tensors): + queue.put((ref, rank_tensor)) sent[id(item)] = ref - for queue in self._queues: - queue.put((ref, item)) return ref if isinstance(item, dict): return {key: replace(child) for key, child in item.items()} @@ -108,30 +200,8 @@ def receive(self, value: Any, *, rank: int, device: str | torch.device) -> Any: """Resolve tensor refs for one consumer rank onto its local device.""" if not 0 <= rank < self.consumer_world_size: raise ValueError(f"Consumer rank {rank} is outside [0, {self.consumer_world_size})") - queue = self._queues[rank] resolved: dict[WorkerTensorRef, torch.Tensor] = {} - def receive_tensor(expected: WorkerTensorRef) -> torch.Tensor: - while True: - if not queue._reader.poll(self.timeout): - raise TimeoutError( - f"Tensor channel {self.channel_id} timed out receiving transfer {expected.transfer_id}" - ) - received_ref, tensor = queue.get() - received_key = (received_ref.transfer_id, received_ref.tensor_index) - expected_key = (expected.transfer_id, expected.tensor_index) - if received_key < expected_key: - # The parent dropped this earlier artifact, normally after - # cancellation. Releasing it here keeps the FIFO usable. - del tensor - continue - if received_ref != expected: - raise RuntimeError( - f"Tensor channel {self.channel_id} expected {expected}, received {received_ref}; " - "consumer calls must preserve producer order" - ) - return tensor - def replace(item: Any) -> Any: if isinstance(item, WorkerTensorRef): if item.channel_id != self.channel_id: @@ -139,17 +209,30 @@ def replace(item: Any) -> Any: cached = resolved.get(item) if cached is not None: return cached - tensor = receive_tensor(item) + received = self._receive_payload(item, rank=rank) + target = torch.device(device) + if isinstance(received, _CudaIpcPayload): + tensor = self._copy_cuda_payload(received, rank=rank, target=target) + else: + tensor = received + expected_shape = self._rank_shape(item.shape, item.shard_dim, rank) + expected_nbytes = item.nbytes + if item.shard_dim is not None: + expected_nbytes = tensor.element_size() + for size in expected_shape: + expected_nbytes *= size if ( - tuple(tensor.shape) != item.shape + tuple(tensor.shape) != expected_shape or str(tensor.dtype) != item.dtype - or str(tensor.device) != item.source_device - or tensor.numel() * tensor.element_size() != item.nbytes + or tensor.numel() * tensor.element_size() != expected_nbytes ): raise RuntimeError( f"Tensor channel {self.channel_id} received incompatible tensor metadata for {item}" ) - target = torch.device(device) + if not isinstance(received, _CudaIpcPayload) and str(tensor.device) != item.source_device: + raise RuntimeError( + f"Tensor channel {self.channel_id} received incompatible tensor metadata for {item}" + ) if tensor.device != target: tensor = tensor.to(target, non_blocking=True) resolved[item] = tensor @@ -164,6 +247,309 @@ def replace(item: Any) -> Any: return replace(value) + def _normalize_shard_dim(self, tensor: torch.Tensor) -> int | None: + if self.shard_dim is None: + return None + if tensor.ndim == 0: + raise ValueError("Tensor channel cannot shard a scalar tensor") + shard_dim = self.shard_dim if self.shard_dim >= 0 else tensor.ndim + self.shard_dim + if not 0 <= shard_dim < tensor.ndim: + raise ValueError(f"Tensor channel shard_dim={self.shard_dim} is invalid for shape {tuple(tensor.shape)}") + if tensor.shape[shard_dim] < self.consumer_world_size: + raise ValueError( + f"Tensor channel cannot shard shape {tuple(tensor.shape)} along dimension {shard_dim} " + f"across {self.consumer_world_size} consumers" + ) + return shard_dim + + def _rank_shape(self, shape: tuple[int, ...], shard_dim: int | None, rank: int) -> tuple[int, ...]: + if shard_dim is None: + return shape + base_size, remainder = divmod(shape[shard_dim], self.consumer_world_size) + rank_shape = list(shape) + rank_shape[shard_dim] = base_size + int(rank < remainder) + return tuple(rank_shape) + + def _create_cuda_pool(self, ref: WorkerTensorRef, tensor: torch.Tensor) -> _CudaIpcPool: + with torch.cuda.device(tensor.device): + storage = torch.empty( + (self.cuda_ipc_slots, *tensor.shape), + dtype=tensor.dtype, + device=tensor.device, + ) + slots = [] + for index in range(self.cuda_ipc_slots): + ready_event = torch.cuda.Event(interprocess=True) + slots.append( + _CudaIpcSlot( + tensor=storage[index], + pending_ranks=set(), + completion_ranks=set(), + ready_event=ready_event, + ready_event_handle=ready_event.ipc_handle(), + ) + ) + pool = _CudaIpcPool( + profile_index=len(self._cuda_pools), + storage=storage, + handle=storage.untyped_storage()._share_cuda_(), + slots=slots, + ) + key = (ref.tensor_index, ref.shape, tensor.dtype, str(tensor.device)) + self._cuda_pools[key] = pool + return pool + + def _ack_index(self, profile_index: int, slot_id: int, rank: int) -> int: + return (profile_index * self.cuda_ipc_slots + slot_id) * self.consumer_world_size + rank + + def _drain_acks(self) -> None: + for pool in self._cuda_pools.values(): + for slot_id, slot in enumerate(pool.slots): + if slot.transfer_key is None: + continue + generation = slot.transfer_key[0] + 1 + acknowledgements = { + rank: self._ack_generations[self._ack_index(pool.profile_index, slot_id, rank)] + for rank in slot.pending_ranks + } + slot.completion_ranks.update( + rank for rank, acknowledged in acknowledgements.items() if acknowledged == generation + ) + slot.pending_ranks = { + rank for rank, acknowledged in acknowledgements.items() if abs(acknowledged) != generation + } + if not slot.pending_ranks: + slot.transfer_key = None + + def _acquire_cuda_slot(self, pool: _CudaIpcPool) -> tuple[int, _CudaIpcSlot]: + deadline = None + while True: + self._drain_acks() + for offset in range(len(pool.slots)): + slot_id = (pool.next_slot_id + offset) % len(pool.slots) + slot = pool.slots[slot_id] + if not slot.pending_ranks: + pool.next_slot_id = (slot_id + 1) % len(pool.slots) + return slot_id, slot + if deadline is None: + deadline = time.monotonic() + self.timeout + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"Tensor channel {self.channel_id} timed out waiting for a reusable CUDA IPC slot") + time.sleep(min(0.00005, remaining)) + + def _stage_cuda_tensor(self, ref: WorkerTensorRef, tensor: torch.Tensor) -> tuple[_CudaIpcPayload, ...] | None: + key = (ref.tensor_index, ref.shape, tensor.dtype, str(tensor.device)) + pool = self._cuda_pools.get(key) + if pool is None: + if len(self._cuda_pools) >= _MAX_CUDA_IPC_POOL_PROFILES: + return None + try: + pool = self._create_cuda_pool(ref, tensor) + except torch.OutOfMemoryError: + return None + slot_id, slot = self._acquire_cuda_slot(pool) + with torch.cuda.device(tensor.device): + producer_stream = torch.cuda.current_stream(tensor.device) + self._wait_for_cuda_slot_reuse( + pool, + slot_id, + slot, + device=tensor.device, + producer_stream=producer_stream, + ) + slot.tensor.copy_(tensor, non_blocking=True) + slot.ready_event.record(producer_stream) + rank_views = ( + (slot.tensor,) * self.consumer_world_size + if ref.shard_dim is None + else torch.tensor_split(slot.tensor, self.consumer_world_size, dim=ref.shard_dim) + ) + slot.pending_ranks = set(range(self.consumer_world_size)) + slot.completion_ranks.clear() + slot.transfer_key = (ref.transfer_id, ref.tensor_index) + slot.was_used = True + return tuple( + _CudaIpcPayload( + ref=ref, + profile_index=pool.profile_index, + slot_id=slot_id, + handle=pool.handle, + ready_event_handle=slot.ready_event_handle, + byte_offset=view.storage_offset() * view.element_size(), + shape=tuple(view.shape), + stride=tuple(view.stride()), + dtype=view.dtype, + ) + for view in rank_views + ) + + def _wait_for_cuda_slot_reuse( + self, + pool: _CudaIpcPool, + slot_id: int, + slot: _CudaIpcSlot, + *, + device: torch.device, + producer_stream: torch.cuda.Stream, + ) -> None: + if not slot.was_used: + return + if device.index is None: + raise RuntimeError("CUDA IPC pool device must have an explicit index") + deadline = time.monotonic() + self.timeout + for rank in sorted(slot.completion_ranks): + completion_queue = self._completion_queues[rank] + handle_key = (pool.profile_index, slot_id, rank) + while handle_key not in self._cuda_completion_handles: + remaining = deadline - time.monotonic() + if remaining <= 0 or not completion_queue._reader.poll(remaining): + raise TimeoutError( + f"Tensor channel {self.channel_id} timed out waiting for CUDA completion metadata" + ) + profile_index, completed_slot_id, event_handle = completion_queue.get() + self._cuda_completion_handles[(profile_index, completed_slot_id, rank)] = event_handle + event_key = (*handle_key, device.index) + completion_event = self._cuda_producer_completion_events.get(event_key) + if completion_event is None: + completion_event = torch.cuda.Event.from_ipc_handle( + device, + self._cuda_completion_handles[handle_key], + ) + self._cuda_producer_completion_events[event_key] = completion_event + if not completion_event.query(): + producer_stream.wait_event(completion_event) + + def _publish_cuda_completion( + self, + payload: _CudaIpcPayload, + *, + rank: int, + device: torch.device, + copy_stream: torch.cuda.Stream, + ) -> None: + event_key = (payload.profile_index, payload.slot_id) + completion_event = self._cuda_consumer_completion_events.get(event_key) + is_new_event = completion_event is None + if completion_event is None: + completion_event = torch.cuda.Event(interprocess=True) + self._cuda_consumer_completion_events[event_key] = completion_event + completion_event.record(copy_stream) + if is_new_event: + self._completion_queues[rank].put((payload.profile_index, payload.slot_id, completion_event.ipc_handle())) + self._ack_cuda_payload(payload, rank=rank) + + def _copy_cuda_payload( + self, + payload: _CudaIpcPayload, + *, + rank: int, + target: torch.device, + ) -> torch.Tensor: + source = torch.device(payload.ref.source_device) + rebuild_device = target if target.type == "cuda" else source + if rebuild_device.index is None: + rebuild_device = torch.device(rebuild_device.type, torch.cuda.current_device()) + handle_key = (payload.handle[1], rebuild_device.index) + storage = self._cuda_storage_cache.get(handle_key) + if storage is None: + redirected_handle = (rebuild_device.index, *payload.handle[1:]) + with torch.cuda.device(rebuild_device): + storage = torch.UntypedStorage._new_shared_cuda(*redirected_handle) + self._cuda_storage_cache[handle_key] = storage + with torch.cuda.device(rebuild_device): + copy_stream = torch.cuda.current_stream(rebuild_device) + event_key = (payload.ready_event_handle, rebuild_device.index) + ready_event = self._cuda_event_cache.get(event_key) + if ready_event is None: + ready_event = torch.cuda.Event.from_ipc_handle(rebuild_device, payload.ready_event_handle) + self._cuda_event_cache[event_key] = ready_event + if not ready_event.query(): + copy_stream.wait_event(ready_event) + mapped = torch.empty(0, dtype=payload.dtype, device=rebuild_device).set_( + storage, + storage_offset=payload.byte_offset // torch.empty((), dtype=payload.dtype).element_size(), + size=payload.shape, + stride=payload.stride, + ) + if target.type == "cuda": + output = torch.empty(payload.shape, dtype=payload.dtype, device=target) + output.copy_(mapped, non_blocking=True) + else: + output = mapped.to(target) + self._publish_cuda_completion(payload, rank=rank, device=rebuild_device, copy_stream=copy_stream) + return output + + def _ack_cuda_payload(self, payload: _CudaIpcPayload, *, rank: int, copied: bool = True) -> None: + generation = payload.ref.transfer_id + 1 + self._ack_generations[self._ack_index(payload.profile_index, payload.slot_id, rank)] = ( + generation if copied else -generation + ) + + def release_local_cuda_ipc(self) -> None: + """Release process-local CUDA IPC mappings after pending work is synchronized.""" + self._cuda_consumer_completion_events.clear() + self._cuda_completion_handles.clear() + self._cuda_producer_completion_events.clear() + self._cuda_event_cache.clear() + self._cuda_storage_cache.clear() + + def discard(self, value: Any, *, rank: int) -> int: + """Consume and release referenced tensors without materializing a device copy.""" + if not 0 <= rank < self.consumer_world_size: + raise ValueError(f"Consumer rank {rank} is outside [0, {self.consumer_world_size})") + discarded: set[WorkerTensorRef] = set() + + def visit(item: Any) -> None: + if isinstance(item, WorkerTensorRef): + if item.channel_id != self.channel_id or item in discarded: + return + received = self._receive_payload(item, rank=rank) + if isinstance(received, _CudaIpcPayload): + self._ack_cuda_payload(received, rank=rank, copied=False) + else: + del received + discarded.add(item) + return + if isinstance(item, dict): + for child in item.values(): + visit(child) + elif isinstance(item, tuple | list): + for child in item: + visit(child) + + visit(value) + return len(discarded) + + def _receive_payload(self, expected: WorkerTensorRef, *, rank: int) -> torch.Tensor | _CudaIpcPayload: + if expected._cuda_payloads is not None: + if len(expected._cuda_payloads) != self.consumer_world_size: + raise RuntimeError(f"Tensor channel {self.channel_id} received incomplete CUDA IPC payload metadata") + payload = expected._cuda_payloads[rank] + if payload.ref != expected: + raise RuntimeError(f"Tensor channel {self.channel_id} received mismatched CUDA IPC payload metadata") + return payload + queue = self._queues[rank] + while True: + if not queue._reader.poll(self.timeout): + raise TimeoutError( + f"Tensor channel {self.channel_id} timed out receiving transfer {expected.transfer_id}" + ) + received_ref, tensor = queue.get() + received_key = (received_ref.transfer_id, received_ref.tensor_index) + expected_key = (expected.transfer_id, expected.tensor_index) + if received_key < expected_key: + # The parent dropped this earlier artifact. Releasing it keeps + # the FIFO usable when cancellation cleanup was delayed. + del tensor + continue + if received_ref != expected: + raise RuntimeError( + f"Tensor channel {self.channel_id} expected {expected}, received {received_ref}; " + "consumer calls must preserve producer order" + ) + return tensor + def contains_ref(self, value: Any) -> bool: """Return whether a nested value contains a ref owned by this channel.""" if isinstance(value, WorkerTensorRef): @@ -179,5 +565,8 @@ def close(self) -> None: if self._closed: return self._closed = True - for queue in self._queues: - queue.close() + self.release_local_cuda_ipc() + for tensor_queue in self._queues: + tensor_queue.close() + for completion_queue in self._completion_queues: + completion_queue.close() diff --git a/tests/integration/test_collectives.py b/tests/integration/test_collectives.py new file mode 100644 index 0000000..43a7ea0 --- /dev/null +++ b/tests/integration/test_collectives.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +import pytest +import torch +import torch.distributed as dist + +from telefuser.core.base_stage import BaseStage +from telefuser.core.config import ModelRuntimeConfig, ParallelConfig +from telefuser.distributed.collectives import all_gather_cat, all_gather_stacked, all_reduce_sum_ +from telefuser.worker import ParallelWorker + + +class _CollectiveStage(BaseStage): + def __init__(self) -> None: + super().__init__( + "collective-integration", + ModelRuntimeConfig( + device_type="cuda", + torch_dtype=torch.float32, + parallel_config=ParallelConfig(device_ids=[0, 1], sp_ulysses_degree=2), + ), + ) + self.empty_cache_after_call = False + + def run(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + rank = dist.get_rank() + local = torch.full((2, 3), float(rank), device=self.device) + stacked = all_gather_stacked(local) + concatenated = all_gather_cat(local, dim=1) + value = torch.tensor([rank + 1.0], device=self.device) + weight = torch.tensor([2.0 * rank + 1.0], device=self.device) + all_reduce_sum_((value, weight)) + return stacked.cpu(), concatenated.cpu(), value.cpu(), weight.cpu() + + +@pytest.mark.distributed +@pytest.mark.gpu +@pytest.mark.multi_gpu +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") +def test_shared_collectives_preserve_rank_order_and_reduce_all_tensors() -> None: + worker = ParallelWorker(_CollectiveStage()) + try: + stacked, concatenated, value, weight = worker.run(sync=True) + finally: + worker.close() + + assert stacked.shape == (2, 2, 3) + torch.testing.assert_close(stacked[0], torch.zeros(2, 3)) + torch.testing.assert_close(stacked[1], torch.ones(2, 3)) + torch.testing.assert_close(concatenated, torch.tensor([[0, 0, 0, 1, 1, 1]]).expand(2, -1).float()) + torch.testing.assert_close(value, torch.tensor([3.0])) + torch.testing.assert_close(weight, torch.tensor([4.0])) diff --git a/tests/integration/test_wan_video_vae_spatial.py b/tests/integration/test_wan_video_vae_spatial.py new file mode 100644 index 0000000..3b2c439 --- /dev/null +++ b/tests/integration/test_wan_video_vae_spatial.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +from copy import deepcopy +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist + +from telefuser.core.base_stage import BaseStage +from telefuser.core.config import ModelRuntimeConfig, ParallelConfig +from telefuser.distributed.vae_spatial import _split_height +from telefuser.models.wan_video_vae import ( + VideoVAE, + _convert_conv3d_to_channels_last_3d, + _count_conv3d, + _enable_spatial_parallel_decode, +) +from telefuser.worker import ParallelWorker + + +class _SpatialVAEParityStage(BaseStage): + def __init__(self) -> None: + super().__init__( + "spatial-vae-parity", + ModelRuntimeConfig( + device_type="cuda", + torch_dtype=torch.float32, + parallel_config=ParallelConfig(device_ids=[0, 1], sp_ulysses_degree=2), + ), + ) + torch.manual_seed(17) + source = VideoVAE(dim=8, z_dim=4).eval() + self.dense_conv2 = deepcopy(source.conv2) + self.dense_decoder = deepcopy(source.decoder) + self.spatial_conv2 = deepcopy(source.conv2) + self.spatial_decoder = deepcopy(source.decoder) + self.empty_cache_after_call = False + + def parallel_models(self) -> None: + self.dense_conv2 = self.dense_conv2.to(self.device) + self.dense_decoder = self.dense_decoder.to(self.device) + self.spatial_conv2 = self.spatial_conv2.to(self.device) + self.spatial_decoder = self.spatial_decoder.to(self.device) + _convert_conv3d_to_channels_last_3d(self.dense_decoder) + _convert_conv3d_to_channels_last_3d(self.spatial_decoder) + vae = SimpleNamespace( + model=SimpleNamespace(decoder=self.spatial_decoder), + parallelism=1, + ) + _enable_spatial_parallel_decode(vae) + + @staticmethod + def _decode_chunks( + conv2: torch.nn.Module, + decoder: torch.nn.Module, + chunks: list[torch.Tensor], + *, + pre_shard: bool = False, + ) -> list[torch.Tensor]: + cache: list[object] = [None] * _count_conv3d(decoder) + outputs = [] + for chunk in chunks: + cache_index = [0] + encoded = conv2(chunk) + kwargs = {} + if pre_shard: + kwargs["input_global_height"] = encoded.shape[-2] + encoded = _split_height(encoded) + outputs.append(decoder(encoded, feat_cache=cache, feat_idx=cache_index, **kwargs).cpu()) + return outputs + + def compare(self, chunks: list[torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]: + dense = self._decode_chunks(self.dense_conv2, self.dense_decoder, chunks) if dist.get_rank() == 0 else [] + spatial = self._decode_chunks(self.spatial_conv2, self.spatial_decoder, chunks) + pre_sharded = self._decode_chunks(self.spatial_conv2, self.spatial_decoder, chunks, pre_shard=True) + return dense, spatial, pre_sharded + + +@pytest.mark.distributed +@pytest.mark.gpu +@pytest.mark.multi_gpu +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") +def test_streaming_spatial_vae_matches_dense_decode_across_causal_chunks() -> None: + torch.manual_seed(29) + chunks = [ + torch.randn(1, 4, 1, 5, 4, dtype=torch.float32), + torch.randn(1, 4, 1, 5, 4, dtype=torch.float32), + ] + worker = ParallelWorker(_SpatialVAEParityStage()) + try: + dense_outputs, spatial_outputs, pre_sharded_outputs = worker.compare(chunks, sync=True) + finally: + worker.close() + + assert len(dense_outputs) == len(spatial_outputs) == len(pre_sharded_outputs) == 2 + for dense, spatial, pre_sharded in zip(dense_outputs, spatial_outputs, pre_sharded_outputs): + torch.testing.assert_close(spatial, dense, rtol=2e-4, atol=2e-4) + torch.testing.assert_close(pre_sharded, dense, rtol=2e-4, atol=2e-4) diff --git a/tests/integration/test_worker_tensor_channel.py b/tests/integration/test_worker_tensor_channel.py index caee279..1e5fb26 100644 --- a/tests/integration/test_worker_tensor_channel.py +++ b/tests/integration/test_worker_tensor_channel.py @@ -1,11 +1,15 @@ from __future__ import annotations +import threading + import pytest import torch import torch.distributed as dist +import torch.multiprocessing as mp from telefuser.core.base_stage import BaseStage from telefuser.core.config import ModelRuntimeConfig, ParallelConfig +from telefuser.distributed.vae_spatial import _gather_height from telefuser.worker import ParallelWorker, WorkerTensorChannel, WorkerTensorRef @@ -41,6 +45,40 @@ def consume(self, tensor: torch.Tensor) -> tuple[str, torch.Tensor]: return str(tensor.device), tensor.cpu() +class _DistributedGPUConsumerStage(BaseStage): + def __init__(self) -> None: + super().__init__( + "distributed-gpu-consumer", + ModelRuntimeConfig( + device_type="cuda", + device_id=2, + parallel_config=ParallelConfig(device_ids=[2, 3], sp_ulysses_degree=2), + ), + ) + + def consume(self, tensor: torch.Tensor) -> tuple[str, torch.Tensor]: + return str(tensor.device), tensor.cpu() + + def gather_height_shards(self, tensor: torch.Tensor) -> tuple[tuple[int, ...], torch.Tensor]: + local_shape = tuple(tensor.shape) + return local_shape, _gather_height(tensor).cpu() + + +class _BlockingRankDiscardChannel(WorkerTensorChannel): + def __init__(self) -> None: + super().__init__(consumer_world_size=2, timeout=30) + context = mp.get_context("spawn") + self.rank_one_started = context.Event() + self.rank_one_release = context.Event() + + def discard(self, value: object, *, rank: int) -> int: + if rank == 1: + self.rank_one_started.set() + if not self.rank_one_release.wait(timeout=30): + raise TimeoutError("Timed out waiting to release consumer rank 1") + return super().discard(value, rank=rank) + + @pytest.mark.distributed @pytest.mark.skipif(torch.cuda.device_count() < 3, reason="requires three CUDA devices") def test_distributed_worker_to_worker_tensor_path_bypasses_parent() -> None: @@ -52,6 +90,10 @@ def test_distributed_worker_to_worker_tensor_path_bypasses_parent() -> None: ) consumer = ParallelWorker(_GPUConsumerStage(), tensor_input_channels=(channel,)) try: + abandoned = producer.reduce(torch.arange(4, dtype=torch.float32), sync=True) + assert isinstance(abandoned, WorkerTensorRef) + assert consumer.discard_tensor_refs(abandoned, sync=True) == 1 + ref = producer.reduce(torch.arange(4, dtype=torch.float32), sync=True) assert isinstance(ref, WorkerTensorRef) assert ref.source_device == "cuda:0" @@ -59,6 +101,70 @@ def test_distributed_worker_to_worker_tensor_path_bypasses_parent() -> None: assert device == "cuda:2" torch.testing.assert_close(result, 2 * torch.arange(4, dtype=torch.float32) + 1) finally: + consumer.close() producer.close() + channel.close() + + +@pytest.mark.distributed +@pytest.mark.skipif(torch.cuda.device_count() < 4, reason="requires four CUDA devices") +def test_discard_waits_for_every_consumer_rank() -> None: + channel = _BlockingRankDiscardChannel() + producer = ParallelWorker( + _DistributedProducerStage(), + tensor_output_channel=channel, + tensor_output_methods=("reduce",), + ) + consumer = ParallelWorker(_DistributedGPUConsumerStage(), tensor_input_channels=(channel,)) + discarded: list[int] = [] + + def discard() -> None: + discarded.append(consumer.discard_tensor_refs(abandoned, sync=True)) + + try: + abandoned = producer.reduce(torch.arange(4, dtype=torch.float32), sync=True) + thread = threading.Thread(target=discard) + thread.start() + assert channel.rank_one_started.wait(timeout=10) + thread.join(timeout=0.1) + assert thread.is_alive() + + channel.rank_one_release.set() + thread.join(timeout=30) + assert not thread.is_alive() + assert discarded == [1] + + current = producer.reduce(torch.arange(4, dtype=torch.float32), sync=True) + device, result = consumer.consume(current, sync=True) + assert device == "cuda:2" + torch.testing.assert_close(result, 2 * torch.arange(4, dtype=torch.float32) + 1) + finally: + channel.rank_one_release.set() + consumer.close() + producer.close() + channel.close() + + +@pytest.mark.distributed +@pytest.mark.skipif(torch.cuda.device_count() < 4, reason="requires four CUDA devices") +def test_sharded_channel_copies_only_each_consumer_rank_height_slice() -> None: + channel = WorkerTensorChannel(consumer_world_size=2, timeout=30, shard_dim=-2) + producer = ParallelWorker( + _DistributedProducerStage(), + tensor_output_channel=channel, + tensor_output_methods=("reduce",), + ) + consumer = ParallelWorker(_DistributedGPUConsumerStage(), tensor_input_channels=(channel,)) + source = torch.arange(30, dtype=torch.float32).reshape(1, 1, 1, 5, 6) + try: + ref = producer.reduce(source, sync=True) + assert isinstance(ref, WorkerTensorRef) + assert ref.shard_dim == 3 + local_shape, gathered = consumer.gather_height_shards(ref, sync=True) + finally: consumer.close() + producer.close() channel.close() + + assert local_shape == (1, 1, 1, 3, 6) + torch.testing.assert_close(gathered, 2 * source + 1) diff --git a/tests/unit/distributed/test_collectives.py b/tests/unit/distributed/test_collectives.py new file mode 100644 index 0000000..94106ef --- /dev/null +++ b/tests/unit/distributed/test_collectives.py @@ -0,0 +1,61 @@ +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from telefuser.distributed.collectives import all_gather_cat, all_gather_stacked, all_reduce_sum_ + + +def test_all_gather_stacked_uses_one_contiguous_output() -> None: + local = torch.arange(6, dtype=torch.float32).reshape(2, 3) + + def gather(output: torch.Tensor, tensor: torch.Tensor, *, group: object) -> None: + ranks = output.view(2, *tensor.shape) + ranks[0].copy_(tensor) + ranks[1].copy_(tensor + 10) + + with patch("telefuser.distributed.collectives.dist.all_gather_into_tensor", side_effect=gather) as mocked: + gathered = all_gather_stacked(local, group=MagicMock(), world_size=2) + + assert gathered.shape == (2, 2, 3) + torch.testing.assert_close(gathered[0], local) + torch.testing.assert_close(gathered[1], local + 10) + mocked.assert_called_once() + + +def test_all_gather_cat_supports_arbitrary_dimensions() -> None: + local = torch.arange(6, dtype=torch.float32).reshape(2, 3) + moved = local.movedim(1, 0).contiguous() + gathered = torch.stack((moved, moved + 10)) + + with patch("telefuser.distributed.collectives.all_gather_stacked", return_value=gathered): + merged = all_gather_cat(local, dim=1, group=MagicMock(), world_size=2) + + torch.testing.assert_close(merged, torch.cat((local, local + 10), dim=1)) + assert merged.is_contiguous() + + +def test_all_gather_helpers_are_no_ops_for_one_rank() -> None: + tensor = torch.randn(2, 3) + + assert all_gather_stacked(tensor, world_size=1).data_ptr() == tensor.data_ptr() + assert all_gather_cat(tensor, dim=-1, world_size=1) is tensor + + +@pytest.mark.parametrize("dim", [-3, 2]) +def test_all_gather_cat_rejects_invalid_dimension(dim: int) -> None: + with pytest.raises(ValueError, match="invalid"): + all_gather_cat(torch.zeros(2, 3), dim=dim, world_size=2) + + +def test_all_reduce_sum_submits_before_waiting() -> None: + tensors = (torch.ones(1), torch.ones(1)) + works = (MagicMock(), MagicMock()) + + with patch("telefuser.distributed.collectives.dist.all_reduce", side_effect=works) as mocked: + all_reduce_sum_(tensors, group=MagicMock()) + + assert mocked.call_count == 2 + assert all(call.kwargs["async_op"] is True for call in mocked.call_args_list) + works[0].wait.assert_called_once() + works[1].wait.assert_called_once() diff --git a/tests/unit/distributed/test_parallel_shard.py b/tests/unit/distributed/test_parallel_shard.py index 463f861..7d19166 100644 --- a/tests/unit/distributed/test_parallel_shard.py +++ b/tests/unit/distributed/test_parallel_shard.py @@ -149,7 +149,7 @@ def test_no_op_when_world_size_1( @patch("telefuser.distributed.parallel_shard.get_attention_strategy") @patch("telefuser.distributed.parallel_shard.get_sp_shard_group") @patch("telefuser.distributed.parallel_shard.get_sp_shard_degree") - @patch("telefuser.distributed.parallel_shard.dist.all_gather") + @patch("telefuser.distributed.parallel_shard.all_gather_cat") def test_unshard_multi_gpu( self, mock_all_gather, mock_get_sp_shard_degree, mock_get_sp_shard_group, mock_get_attention_strategy ): @@ -161,11 +161,7 @@ def test_unshard_multi_gpu( mock_get_sp_shard_group.return_value = MagicMock() mock_mesh = MagicMock() - def mock_gather(output_list, input_tensor, group): - for i, t in enumerate(output_list): - t.copy_(input_tensor + i) - - mock_all_gather.side_effect = mock_gather + mock_all_gather.side_effect = lambda tensor, *, dim, **_: torch.cat((tensor, tensor + 1), dim=dim) tensors = [torch.randn(2, 5, 64)] result = sequence_parallel_unshard(mock_mesh, tensors, [1], [10]) @@ -176,7 +172,7 @@ def mock_gather(output_list, input_tensor, group): @patch("telefuser.distributed.parallel_shard.get_attention_strategy") @patch("telefuser.distributed.parallel_shard.get_sp_shard_group") @patch("telefuser.distributed.parallel_shard.get_sp_shard_degree") - @patch("telefuser.distributed.parallel_shard.dist.all_gather") + @patch("telefuser.distributed.parallel_shard.all_gather_cat") def test_unshard_multiple_tensors( self, mock_all_gather, mock_get_sp_shard_degree, mock_get_sp_shard_group, mock_get_attention_strategy ): @@ -188,11 +184,7 @@ def test_unshard_multiple_tensors( mock_get_sp_shard_group.return_value = MagicMock() mock_mesh = MagicMock() - def mock_gather(output_list, input_tensor, group): - for t in output_list: - t.copy_(input_tensor) - - mock_all_gather.side_effect = mock_gather + mock_all_gather.side_effect = lambda tensor, *, dim, **_: torch.cat((tensor, tensor), dim=dim) tensors = [torch.randn(2, 5, 64), torch.randn(3, 8, 32)] result = sequence_parallel_unshard(mock_mesh, tensors, [1, 1], [10, 16]) @@ -289,27 +281,23 @@ def test_no_op_when_world_size_1(self, mock_get_cfg_world_size): @patch("telefuser.distributed.parallel_shard.get_cfg_world_size") @patch("telefuser.distributed.parallel_shard.get_cfg_group") - @patch("telefuser.distributed.parallel_shard.dist.all_gather_into_tensor") + @patch("telefuser.distributed.parallel_shard.all_gather_cat") def test_unshard_gather_into_tensor(self, mock_all_gather, mock_get_cfg_group, mock_get_cfg_world_size): - """Test unsharding uses all_gather_into_tensor.""" + """Test unsharding delegates to the shared gather primitive.""" from telefuser.distributed.parallel_shard import cfg_parallel_unshard mock_get_cfg_world_size.return_value = 2 mock_get_cfg_group.return_value = MagicMock() mock_mesh = MagicMock() - def mock_gather(output, input_tensor, group): - output.fill_(1.0) - - mock_all_gather.side_effect = mock_gather + mock_all_gather.side_effect = lambda tensor, *, dim, **_: torch.cat((tensor, tensor), dim=dim) tensors = [torch.randn(2, 10, 64), torch.randn(3, 8, 32)] result = cfg_parallel_unshard(mock_mesh, tensors) assert len(result) == 2 - # Output shape is (cfg_world_size, *tensor.shape[1:]) - assert result[0].shape == (2, 10, 64) - assert result[1].shape == (2, 8, 32) + assert result[0].shape == (4, 10, 64) + assert result[1].shape == (6, 8, 32) assert mock_all_gather.call_count == 2 @@ -320,7 +308,7 @@ class TestShardUnshardRoundTrip: @patch("telefuser.distributed.parallel_shard.get_sp_shard_degree") @patch("telefuser.distributed.parallel_shard.get_sp_shard_rank") @patch("telefuser.distributed.parallel_shard.get_sp_shard_group") - @patch("telefuser.distributed.parallel_shard.dist.all_gather") + @patch("telefuser.distributed.parallel_shard.all_gather_cat") def test_sequence_parallel_roundtrip( self, mock_all_gather, @@ -338,11 +326,7 @@ def test_sequence_parallel_roundtrip( mock_get_sp_shard_group.return_value = MagicMock() mock_mesh = MagicMock() - def mock_gather(output_list, input_tensor, group): - for i, t in enumerate(output_list): - t.copy_(input_tensor + i * 0.1) # Slight difference per rank - - mock_all_gather.side_effect = mock_gather + mock_all_gather.side_effect = lambda tensor, *, dim, **_: torch.cat((tensor, tensor + 0.1), dim=dim) original_shape = (2, 10, 64) tensor = torch.randn(*original_shape) diff --git a/tests/unit/models/test_lingbot_video_dit.py b/tests/unit/models/test_lingbot_video_dit.py index 255edfb..c8430ab 100644 --- a/tests/unit/models/test_lingbot_video_dit.py +++ b/tests/unit/models/test_lingbot_video_dit.py @@ -1,5 +1,7 @@ from __future__ import annotations +from unittest.mock import patch + import pytest import torch import torch.nn.functional as F @@ -52,6 +54,49 @@ def test_attention_dispatcher_sdpa_preserves_native_attention_result() -> None: assert torch.equal(actual, expected) +def test_ulysses_submits_all_qkv_collectives_before_waiting() -> None: + module = LingBotVideoAttention(hidden_size=8, num_heads=2, norm_eps=1e-6, qkv_bias=True, out_bias=True) + module.set_attention_config(AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA)) + module.set_ulysses_group(object()) + hidden_states = torch.randn(1, 3, 8) + rotary = torch.ones(3, 2, dtype=torch.complex64) + events: list[str] = [] + submit_index = 0 + + def submit(tensor: torch.Tensor, group: object): + nonlocal submit_index + del group + name = ("q", "k", "v")[submit_index] + submit_index += 1 + events.append(f"submit-{name}") + + def wait() -> torch.Tensor: + events.append(f"wait-{name}") + return tensor + + return wait + + def gather(tensor: torch.Tensor, group: object, *, num_heads: int): + del group, num_heads + events.append("submit-output") + + def wait() -> torch.Tensor: + events.append("wait-output") + return tensor + + return wait + + with ( + patch("telefuser.models.lingbot_video_dit.dist.is_initialized", return_value=True), + patch("telefuser.models.lingbot_video_dit.dist.get_world_size", return_value=2), + patch("telefuser.models.lingbot_video_dit.ulysses_scatter_heads", side_effect=submit), + patch("telefuser.models.lingbot_video_dit.ulysses_gather_heads", side_effect=gather), + ): + module(hidden_states, rotary) + + assert events == ["submit-q", "submit-k", "submit-v", "wait-q", "wait-k", "wait-v", "submit-output", "wait-output"] + + def test_transformer_rejects_non_patch_aligned_latent_geometry() -> None: model = LingBotVideoTransformer3DModel( in_channels=16, diff --git a/tests/unit/pipelines/lingbot_world_fast/test_module_loading.py b/tests/unit/pipelines/lingbot_world_fast/test_module_loading.py index 323db05..024a980 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_module_loading.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_module_loading.py @@ -97,7 +97,7 @@ def test_non_colocated_dit_and_vae_workers_use_direct_tensor_channel(tmp_path) - assert channel_cls.call_args_list[0].args == (2,) assert channel_cls.call_args_list[0].kwargs == {"timeout": 600} assert channel_cls.call_args_list[1].args == (1,) - assert channel_cls.call_args_list[1].kwargs == {"timeout": 600} + assert channel_cls.call_args_list[1].kwargs == {"timeout": 600, "shard_dim": None} assert pipeline.uses_direct_condition_handoff is True assert pipeline.uses_direct_vae_handoff is True assert worker_cls.call_args_list[0].kwargs == { diff --git a/tests/unit/pipelines/lingbot_world_fast/test_streaming.py b/tests/unit/pipelines/lingbot_world_fast/test_streaming.py index 955d6ca..1ae51c0 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_streaming.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_streaming.py @@ -2,7 +2,7 @@ import threading import time -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import torch @@ -16,6 +16,7 @@ LingBotWorldFastSessionStatus, ) from telefuser.pipelines.lingbot_world_fast.streaming import LingBotWorldFastStreamingRuntime +from telefuser.worker.tensor_channel import WorkerTensorRef class _Worker: @@ -25,6 +26,7 @@ def __init__(self, name: str, release_order: list[str]): self.calls: list[dict[str, object]] = [] self.release_calls: list[int] = [] self.release_failures = 0 + self.discard_calls: list[WorkerTensorRef] = [] self.closed = False def encode_condition_chunk(self, **kwargs): @@ -44,6 +46,11 @@ def release_cache(self, cache_handle: int, sync: bool = False): self.release_order.append(self.name) return True + def discard_tensor_refs(self, ref: WorkerTensorRef, sync: bool = False) -> int: + assert sync + self.discard_calls.append(ref) + return 1 + def close(self): self.closed = True @@ -55,6 +62,7 @@ def __init__(self, release_order: list[str]): self.calls: list[dict[str, object]] = [] self.inference_modes: list[bool] = [] self.release_calls: list[int] = [] + self.discard_calls: list[object] = [] self.fail_denoise = False self.started: threading.Event | None = None self.release: threading.Event | None = None @@ -73,11 +81,18 @@ def denoise_and_update_cache(self, **kwargs): def advance_noise(self, cache_handle: int): self.advance_calls.append(cache_handle) - def release_cache(self, cache_handle: int): + def release_cache(self, cache_handle: int, sync: bool | None = None): + if sync is not None: + assert sync self.release_calls.append(cache_handle) self.release_order.append("denoise") return True + def discard_tensor_refs(self, value: object, sync: bool = False) -> int: + assert sync + self.discard_calls.append(value) + return len(LingBotWorldFastStreamingRuntime._direct_transfer_refs(value)) + class _Pipeline: device = "cpu" @@ -185,6 +200,74 @@ def test_streaming_session_routes_one_chunk_through_three_stages() -> None: assert runtime.cache_handle is None +def test_direct_latent_cancellation_drains_only_the_cancelled_fifo_prefix() -> None: + pipeline = _Pipeline() + streaming_runtime = LingBotWorldFastStreamingRuntime(pipeline) + first = WorkerTensorRef("channel", 0, 0, (1,), "torch.float32", "cuda:0", 4) + second = WorkerTensorRef("channel", 1, 0, (1,), "torch.float32", "cuda:0", 4) + streaming_runtime._track_direct_transfer(streaming_runtime._direct_latent_transfers, "first", first) + streaming_runtime._track_direct_transfer(streaming_runtime._direct_latent_transfers, "second", second) + try: + with patch("telefuser.pipelines.lingbot_world_fast.streaming.ParallelWorker", _Worker): + streaming_runtime._cancel_direct_transfers( + streaming_runtime._direct_latent_transfers, + "second", + pipeline.vae_decode_worker, + "latent", + ) + assert pipeline.vae_decode_worker.discard_calls == [] + + streaming_runtime._consume_direct_transfer( + streaming_runtime._direct_latent_transfers, + "first", + first, + pipeline.vae_decode_worker, + "latent", + ) + assert pipeline.vae_decode_worker.discard_calls == [second] + finally: + streaming_runtime.close() + + +def test_close_session_discards_its_unconsumed_direct_latent() -> 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=1, + height=8, + width=8, + frame_tokens=1, + chunk_size=1, + max_attention_size=1, + cache_handle=21, + ) + streaming_runtime = LingBotWorldFastStreamingRuntime(pipeline) + session = streaming_runtime.create_session(runtime) + condition_ref = WorkerTensorRef("condition", 0, 0, (1,), "torch.float32", "cuda:0", 4) + condition = {"latent_condition": condition_ref} + latent_ref = WorkerTensorRef("latent", 0, 0, (1,), "torch.float32", "cuda:0", 4) + streaming_runtime._track_direct_transfer( + streaming_runtime._direct_condition_transfers, + session.session_id, + condition, + ) + streaming_runtime._track_direct_transfer( + streaming_runtime._direct_latent_transfers, + session.session_id, + latent_ref, + ) + try: + with patch("telefuser.pipelines.lingbot_world_fast.streaming.ParallelWorker", (_Worker, _Denoise)): + streaming_runtime.close_session(session) + assert pipeline.denoise_stage.discard_calls == [condition] + assert pipeline.vae_decode_worker.discard_calls == [latent_ref] + finally: + streaming_runtime.close() + + def test_streaming_session_prefetches_two_conditions_ahead_of_controls() -> None: pipeline = _Pipeline() runtime = LingBotWorldFastGenerationSession( diff --git a/tests/unit/worker/test_ray_worker.py b/tests/unit/worker/test_ray_worker.py new file mode 100644 index 0000000..b349a05 --- /dev/null +++ b/tests/unit/worker/test_ray_worker.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +import os +from unittest.mock import patch + +import pytest + +from telefuser.core.config import RayConfig, RayGPUConfig +from telefuser.worker.ray_worker import RayWorker + + +def _ray_worker(*, num_gpus: int, memory_limit: float = 0.0) -> RayWorker: + worker = RayWorker.__new__(RayWorker) + worker.worker_id = "test-stage" + worker.ray_config = RayConfig( + gpu_config=RayGPUConfig(num_gpus=num_gpus, memory_limit=memory_limit), + memory_gb=0, + ) + return worker + + +def test_setup_resources_preserves_ray_visible_devices(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,7") + worker = _ray_worker(num_gpus=2, memory_limit=0.5) + + with ( + patch("telefuser.worker.ray_worker.current_platform.device_count", return_value=2), + patch("telefuser.worker.ray_worker.current_platform.set_device") as set_device, + patch("telefuser.worker.ray_worker.torch.cuda.set_per_process_memory_fraction") as set_fraction, + ): + worker._setup_resources() + + assert os.environ["CUDA_VISIBLE_DEVICES"] == "4,7" + set_device.assert_called_once_with(0) + set_fraction.assert_called_once_with(0.5, device=0) + + +def test_setup_resources_rejects_insufficient_ray_gpu_assignment() -> None: + worker = _ray_worker(num_gpus=2) + + with ( + patch("telefuser.worker.ray_worker.current_platform.device_count", return_value=1), + patch("telefuser.worker.ray_worker.current_platform.set_device") as set_device, + pytest.raises(RuntimeError, match="assigned 1 visible GPUs"), + ): + worker._setup_resources() + + set_device.assert_not_called() diff --git a/tests/unit/worker/test_tensor_channel.py b/tests/unit/worker/test_tensor_channel.py index 71fd17f..727d9c9 100644 --- a/tests/unit/worker/test_tensor_channel.py +++ b/tests/unit/worker/test_tensor_channel.py @@ -46,17 +46,25 @@ def _spawn_receive(channel, metadata_queue, result_queue, release_event) -> None def _spawn_send_cuda(channel, metadata_queue, release_event) -> None: torch.cuda.set_device(0) - tensor = torch.arange(4, dtype=torch.float32, device="cuda:0") - metadata_queue.put(channel.send(tensor)) + for offset in range(4): + tensor = torch.empty(4, dtype=torch.float32, device="cuda:0") + torch.cuda._sleep(50_000_000) + tensor.copy_(torch.arange(4, dtype=torch.float32, device="cuda:0") + offset) + metadata_queue.put(channel.send(tensor)) release_event.wait(timeout=30) def _spawn_receive_cuda(channel, metadata_queue, result_queue, release_event) -> None: torch.cuda.set_device(1) - ref = metadata_queue.get() - tensor = channel.receive(ref, rank=0, device="cuda:1") + outputs = [] + for _ in range(4): + ref = metadata_queue.get() + tensor = channel.receive(ref, rank=0, device="cuda:1") + outputs.append(tensor) torch.cuda.synchronize(1) - result_queue.put((str(tensor.device), tensor.cpu().tolist())) + values = [output.cpu().tolist() for output in outputs] + channel.release_local_cuda_ipc() + result_queue.put((str(tensor.device), values)) release_event.set() @@ -81,6 +89,34 @@ def test_tensor_channel_keeps_parent_artifact_metadata_only_and_fans_out() -> No channel.close() +def test_tensor_channel_sends_only_each_consumers_tensor_shard() -> None: + channel = WorkerTensorChannel(consumer_world_size=2, timeout=1, shard_dim=-2) + source = torch.arange(30, dtype=torch.float32).reshape(1, 1, 1, 5, 6) + try: + artifact = channel.send(source) + rank_zero = channel.receive(artifact, rank=0, device="cpu") + rank_one = channel.receive(artifact, rank=1, device="cpu") + finally: + channel.close() + + assert artifact.shape == source.shape + assert artifact.shard_dim == 3 + assert artifact.nbytes == source.numel() * source.element_size() + assert rank_zero.shape[-2] == 3 + assert rank_one.shape[-2] == 2 + torch.testing.assert_close(torch.cat((rank_zero, rank_one), dim=-2), source) + + +@pytest.mark.parametrize("shape", [(), (1, 2)]) +def test_tensor_channel_rejects_invalid_sharding(shape: tuple[int, ...]) -> None: + channel = WorkerTensorChannel(consumer_world_size=2, timeout=1, shard_dim=-2) + try: + with pytest.raises(ValueError, match="cannot shard|shard_dim"): + channel.send(torch.zeros(shape)) + finally: + channel.close() + + def test_tensor_channel_preserves_nested_container_types() -> None: channel = WorkerTensorChannel(consumer_world_size=1, timeout=1) try: @@ -120,6 +156,19 @@ def test_tensor_channel_discards_cancelled_earlier_transfer() -> None: torch.testing.assert_close(resolved, torch.ones(1)) +def test_tensor_channel_explicitly_discards_terminal_transfer() -> None: + channel = WorkerTensorChannel(consumer_world_size=1, timeout=1) + try: + abandoned = channel.send((torch.zeros(1), {"duplicate": torch.zeros(1)})) + assert channel.discard(abandoned, rank=0) == 2 + current = channel.send(torch.ones(1)) + resolved = channel.receive(current, rank=0, device="cpu") + finally: + channel.close() + + torch.testing.assert_close(resolved, torch.ones(1)) + + def test_tensor_channel_validates_bindings_and_rank() -> None: channel = WorkerTensorChannel(consumer_world_size=2, timeout=1) try: @@ -134,6 +183,9 @@ def test_tensor_channel_validates_bindings_and_rank() -> None: finally: channel.close() + with pytest.raises(ValueError, match="cuda_ipc_slots"): + WorkerTensorChannel(consumer_world_size=1, cuda_ipc_slots=0) + def test_tensor_channel_transfers_between_independent_spawned_processes() -> None: context = mp.get_context("spawn") @@ -177,8 +229,28 @@ def test_parallel_workers_exchange_tensor_without_parent_materialization() -> No result = consumer.consume(ref, sync=True) torch.testing.assert_close(result, torch.arange(4, dtype=torch.float32) + 1) finally: + consumer.close() producer.close() + channel.close() + + +def test_parallel_worker_discards_unconsumed_tensor_refs_on_consumer_ranks() -> None: + channel = WorkerTensorChannel(consumer_world_size=1, timeout=10) + producer = ParallelWorker( + _TensorProducerStage(), + tensor_output_channel=channel, + tensor_output_methods=("produce",), + ) + consumer = ParallelWorker(_TensorConsumerStage(), tensor_input_channels=(channel,)) + try: + abandoned, _ = producer.produce(sync=True) + assert consumer.discard_tensor_refs(abandoned, sync=True) == 1 + current, _ = producer.produce(sync=True) + result = consumer.consume(current, sync=True) + torch.testing.assert_close(result, torch.arange(4, dtype=torch.float32) + 1) + finally: consumer.close() + producer.close() channel.close() @@ -197,9 +269,9 @@ def test_tensor_channel_uses_cuda_ipc_and_cross_gpu_peer_copy() -> None: consumer.start() device, values = result_queue.get() assert device == "cuda:1" - assert values == [0.0, 1.0, 2.0, 3.0] - producer.join(timeout=30) + assert values == [[float(index + offset) for index in range(4)] for offset in range(4)] consumer.join(timeout=30) + producer.join(timeout=30) assert producer.exitcode == 0 assert consumer.exitcode == 0 finally: diff --git a/tools/validation/benchmark_tensor_channel_vs_sglang.py b/tools/validation/benchmark_tensor_channel_vs_sglang.py new file mode 100644 index 0000000..d3fb2d0 --- /dev/null +++ b/tools/validation/benchmark_tensor_channel_vs_sglang.py @@ -0,0 +1,283 @@ +#!/usr/bin/env python3 +"""Compare TeleFuser direct tensor handoff with SGLang's CUDA IPC pool.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import statistics +import sys +import time +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import torch +import torch.multiprocessing as mp + +from telefuser.worker import WorkerTensorChannel + + +def _dtype(name: str) -> torch.dtype: + value = getattr(torch, name, None) + if not isinstance(value, torch.dtype): + raise ValueError(f"Unsupported dtype: {name}") + return value + + +def _telefuser_producer( + channel: WorkerTensorChannel, + commands: Any, + references: Any, + ready: Any, + shape: tuple[int, ...], + dtype_name: str, + device_index: int, +) -> None: + torch.cuda.set_device(device_index) + source = torch.ones(shape, dtype=_dtype(dtype_name), device=f"cuda:{device_index}") + torch.cuda.synchronize(device_index) + ready.put(None) + while commands.get() is not None: + references.put(channel.send(source)) + + +def _telefuser_consumer( + channel: WorkerTensorChannel, + commands: Any, + completed: Any, + ready: Any, + device_index: int, +) -> None: + torch.cuda.set_device(device_index) + torch.empty(1, device=f"cuda:{device_index}") + torch.cuda.synchronize(device_index) + ready.put(None) + try: + while True: + reference = commands.get() + if reference is None: + break + output = channel.receive(reference, rank=0, device=f"cuda:{device_index}") + torch.cuda.synchronize(device_index) + completed.put((tuple(output.shape), output.numel() * output.element_size())) + finally: + channel.release_local_cuda_ipc() + + +def _sglang_producer( + commands: Any, + references: Any, + ready: Any, + shape: tuple[int, ...], + dtype_name: str, + device_index: int, +) -> None: + import sglang.srt.utils.cuda_ipc_transport_utils as transport + + transport.get_server_args = lambda: SimpleNamespace(tp_size=1) + torch.cuda.set_device(device_index) + source = torch.ones(shape, dtype=_dtype(dtype_name), device=f"cuda:{device_index}") + source_bytes = source.view(torch.int8).view(-1) + pool = transport.MmItemMemoryPool(source_bytes.numel(), recycle_interval=60, base_gpu_id=device_index) + torch.cuda.synchronize(device_index) + ready.put(None) + try: + while commands.get() is not None: + with pool._lock: + pool.recycle_chunks() + pool.merge_chunks() + sync_meta, pool_slice, byte_offset = pool.return_a_slice_tensor_with_flag(source) + if pool_slice is None: + raise RuntimeError("SGLang IPC pool did not recycle its single benchmark slot") + pool_slice.copy_(source_bytes, non_blocking=True) + references.put( + transport.CudaIpcTensorTransportProxy( + data=pool_slice, + info_data=source, + sync_buffer_meta=sync_meta, + pool_ipc_handle=pool._pool_ipc_handle, + pool_byte_offset=byte_offset, + pool_device_index=pool._pool_device_index, + ) + ) + finally: + pool.shutdown() + pool.clear_sync_flag_list() + + +def _sglang_consumer(commands: Any, completed: Any, ready: Any, device_index: int) -> None: + from sglang.srt.utils.cuda_ipc_transport_utils import _pool_handle_cache_clear + + torch.cuda.set_device(device_index) + torch.empty(1, device=f"cuda:{device_index}") + torch.cuda.synchronize(device_index) + ready.put(None) + try: + while True: + proxy = commands.get() + if proxy is None: + break + output = proxy.reconstruct_on_target_device(device_index, consumer_count=1) + torch.cuda.synchronize(device_index) + completed.put((tuple(output.shape), output.numel() * output.element_size())) + finally: + _pool_handle_cache_clear() + + +def _run_path( + producer_target: Any, + consumer_target: Any, + *, + shape: tuple[int, ...], + dtype_name: str, + source_device: int, + target_device: int, + warmup: int, + iterations: int, + channel: WorkerTensorChannel | None = None, +) -> list[float]: + context = mp.get_context("spawn") + producer_commands = context.SimpleQueue() + consumer_commands = context.SimpleQueue() + references = context.SimpleQueue() + completed = context.SimpleQueue() + ready = context.SimpleQueue() + common_producer_args = (producer_commands, references, ready, shape, dtype_name, source_device) + producer_args = (channel, *common_producer_args) if channel is not None else common_producer_args + common_consumer_args = (consumer_commands, completed, ready, target_device) + consumer_args = (channel, *common_consumer_args) if channel is not None else common_consumer_args + producer = context.Process(target=producer_target, args=producer_args) + consumer = context.Process(target=consumer_target, args=consumer_args) + producer.start() + consumer.start() + expected_nbytes = math.prod(shape) * torch.empty((), dtype=_dtype(dtype_name)).element_size() + timings = [] + try: + ready.get() + ready.get() + for index in range(warmup + iterations): + started_at = time.perf_counter_ns() + producer_commands.put(True) + reference = references.get() + consumer_commands.put(reference) + output_shape, output_nbytes = completed.get() + elapsed_ms = (time.perf_counter_ns() - started_at) / 1_000_000 + if output_shape != shape or output_nbytes != expected_nbytes: + raise RuntimeError( + f"Transport returned shape={output_shape}, nbytes={output_nbytes}; " + f"expected shape={shape}, nbytes={expected_nbytes}" + ) + if index >= warmup: + timings.append(elapsed_ms) + finally: + consumer_commands.put(None) + consumer.join(timeout=30) + producer_commands.put(None) + producer.join(timeout=30) + for process in (consumer, producer): + if process.is_alive(): + process.terminate() + process.join(timeout=5) + for queue in (producer_commands, consumer_commands, references, completed, ready): + queue.close() + if channel is not None: + channel.close() + if producer.exitcode != 0 or consumer.exitcode != 0: + raise RuntimeError(f"Transport workers failed: producer={producer.exitcode}, consumer={consumer.exitcode}") + return timings + + +def _summary(timings: list[float], nbytes: int, copy_count: int) -> dict[str, float | int]: + ordered = sorted(timings) + p50_ms = statistics.median(ordered) + p95_ms = ordered[math.ceil(0.95 * len(ordered)) - 1] + return { + "p50_ms": round(p50_ms, 4), + "p95_ms": round(p95_ms, 4), + "mean_ms": round(statistics.mean(ordered), 4), + "logical_gib_per_second_at_p50": round(nbytes / (p50_ms / 1000) / (1024**3), 3), + "device_copy_count": copy_count, + "device_bytes_per_transfer": nbytes * copy_count, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--shape", default="1,16,4,60,104") + parser.add_argument("--dtype", default="bfloat16") + parser.add_argument("--source-device", type=int, default=0) + parser.add_argument("--target-device", type=int, default=2) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iterations", type=int, default=200) + parser.add_argument("--max-p50-ratio", type=float, default=1.05) + parser.add_argument("--max-p95-ratio", type=float, default=1.10) + parser.add_argument("--p95-jitter-ms", type=float, default=0.05) + parser.add_argument( + "--sglang-python", + type=Path, + default=Path(__file__).resolve().parents[2] / "work_dirs" / "sglang" / "python", + ) + args = parser.parse_args() + shape = tuple(int(value) for value in args.shape.split(",")) + sglang_python = str(args.sglang_python.resolve()) + sys.path.append(sglang_python) + os.environ["PYTHONPATH"] = os.pathsep.join(filter(None, (sglang_python, os.environ.get("PYTHONPATH")))) + nbytes = math.prod(shape) * torch.empty((), dtype=_dtype(args.dtype)).element_size() + + telefuser_timings = _run_path( + _telefuser_producer, + _telefuser_consumer, + shape=shape, + dtype_name=args.dtype, + source_device=args.source_device, + target_device=args.target_device, + warmup=args.warmup, + iterations=args.iterations, + channel=WorkerTensorChannel(consumer_world_size=1, timeout=30), + ) + sglang_timings = _run_path( + _sglang_producer, + _sglang_consumer, + shape=shape, + dtype_name=args.dtype, + source_device=args.source_device, + target_device=args.target_device, + warmup=args.warmup, + iterations=args.iterations, + ) + telefuser = _summary(telefuser_timings, nbytes, copy_count=2) + sglang = _summary(sglang_timings, nbytes, copy_count=2) + p50_ratio = float(telefuser["p50_ms"]) / float(sglang["p50_ms"]) + p95_ratio = float(telefuser["p95_ms"]) / float(sglang["p95_ms"]) + p95_limit_ms = max( + float(sglang["p95_ms"]) * args.max_p95_ratio, + float(sglang["p95_ms"]) + args.p95_jitter_ms, + ) + result = { + "shape": shape, + "dtype": args.dtype, + "bytes": nbytes, + "source_device": args.source_device, + "target_device": args.target_device, + "iterations": args.iterations, + "telefuser": telefuser, + "sglang_cuda_ipc_pool": sglang, + "telefuser_to_sglang_p50_ratio": round(p50_ratio, 4), + "telefuser_to_sglang_p95_ratio": round(p95_ratio, 4), + "p95_limit_ms": round(p95_limit_ms, 4), + "passes": p50_ratio <= args.max_p50_ratio and float(telefuser["p95_ms"]) <= p95_limit_ms, + } + print(json.dumps(result, indent=2)) + if not result["passes"]: + raise SystemExit( + "TeleFuser latency ratio exceeds the allowed limit: " + f"p50={p50_ratio:.3f}/{args.max_p50_ratio:.3f}, " + f"p95_ms={float(telefuser['p95_ms']):.4f}/{p95_limit_ms:.4f}" + ) + + +if __name__ == "__main__": + main() From 284996dd616cfd44a55523687b7f2a63a281abb9 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Mon, 3 Aug 2026 09:32:27 +0000 Subject: [PATCH 10/11] test(ci): enforce CPU-only local validation Allow the local CI entry point to reuse an existing environment while explicitly hiding CUDA and binding pytest to the active Python interpreter. Make the RayWorker resource test simulate its CUDA branch so it remains deterministic on GitHub CPU runners. Verification:\n- bash scripts/run_ci_tests.sh --skip-install\n- unit tests: 940 passed, 3 skipped, 95 deselected\n- server tests: 62 passed\n- ruff check, format, and import checks --- scripts/run_ci_tests.sh | 34 +++++++++++++++++++++++----- tests/unit/worker/test_ray_worker.py | 7 +++--- 2 files changed, 32 insertions(+), 9 deletions(-) diff --git a/scripts/run_ci_tests.sh b/scripts/run_ci_tests.sh index eb09e28..aedd62a 100755 --- a/scripts/run_ci_tests.sh +++ b/scripts/run_ci_tests.sh @@ -4,6 +4,16 @@ set -e +skip_install=false +if [ "${1:-}" = "--skip-install" ]; then + skip_install=true + shift +fi +if [ "$#" -ne 0 ]; then + echo "Usage: $0 [--skip-install]" + exit 2 +fi + echo "==========================================" echo "Running CI Tests Locally" echo "==========================================" @@ -39,10 +49,15 @@ if [ ! -f "pyproject.toml" ]; then fi # Install dependencies if needed -print_section "Installing dependencies" -pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu -q -pip install -e ".[dev]" -q -check_result "Dependencies installation" +if [ "$skip_install" = true ]; then + print_section "Skipping dependency installation" + echo "Using dependencies from the active Python environment" +else + print_section "Installing dependencies" + pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu -q + pip install -e ".[dev]" -q + check_result "Dependencies installation" +fi # Run lint checks print_section "Running lint checks" @@ -55,16 +70,23 @@ check_result "Ruff format check" ruff check --select I telefuser tests check_result "Import check" +# A CUDA-enabled local PyTorch build is valid for local CI, but the tests below +# must observe the same CPU-only runtime contract as the GitHub runners. +print_section "Enforcing CPU-only test execution" +export CUDA_VISIBLE_DEVICES="" +python -c 'import torch; assert not torch.cuda.is_available(), "CUDA must be hidden during CPU CI tests"' +check_result "CPU-only runtime" + # Run unit tests print_section "Running unit tests" -pytest tests/unit -v \ +python -m pytest tests/unit -v \ -m "not gpu and not distributed and not slow and not quant" \ --tb=short check_result "Unit tests" # Run server pytest tests (includes OpenAI API tests) print_section "Running server pytest tests (includes OpenAI API)" -pytest tests/server/ -v \ +python -m pytest tests/server/ -v \ -m "not gpu and not distributed and not slow" \ --tb=short check_result "Server pytest tests (including OpenAI API)" diff --git a/tests/unit/worker/test_ray_worker.py b/tests/unit/worker/test_ray_worker.py index b349a05..b0a8a30 100644 --- a/tests/unit/worker/test_ray_worker.py +++ b/tests/unit/worker/test_ray_worker.py @@ -24,14 +24,15 @@ def test_setup_resources_preserves_ray_visible_devices(monkeypatch: pytest.Monke worker = _ray_worker(num_gpus=2, memory_limit=0.5) with ( - patch("telefuser.worker.ray_worker.current_platform.device_count", return_value=2), - patch("telefuser.worker.ray_worker.current_platform.set_device") as set_device, + patch("telefuser.worker.ray_worker.current_platform") as current_platform, patch("telefuser.worker.ray_worker.torch.cuda.set_per_process_memory_fraction") as set_fraction, ): + current_platform.device_type = "cuda" + current_platform.device_count.return_value = 2 worker._setup_resources() assert os.environ["CUDA_VISIBLE_DEVICES"] == "4,7" - set_device.assert_called_once_with(0) + current_platform.set_device.assert_called_once_with(0) set_fraction.assert_called_once_with(0.5, device=0) From fec42399cd3da75d4fd0b2d413ce92f9fc71a312 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Mon, 3 Aug 2026 10:01:26 +0000 Subject: [PATCH 11/11] docs(lingbot): document four-GPU real-time performance Publish the validated four-H100 LingBot-World v2 real-time gate in the project and example READMEs. Consolidate the bilingual AIPerf benchmark documentation around the 77-frame compute gate and the current one-minute LiveKit replay, including exact metric boundaries and reproducible commands. Remove stale NEW labels from the LingBot-Video news entries. Verification: AIPerf one-minute stream profile 1/1 succeeded; mkdocs build --strict; git diff --check. --- README.md | 11 ++++- README_zh.md | 9 +++- docs/en/benchmark_aiperf.md | 82 ++++++++++++++++++++++++------------- docs/zh/benchmark_aiperf.md | 74 ++++++++++++++++++++++----------- examples/lingbot/README.md | 37 +++++++++++++++++ 5 files changed, 157 insertions(+), 56 deletions(-) diff --git a/README.md b/README.md index 3ebebdf..b81ab45 100644 --- a/README.md +++ b/README.md @@ -17,9 +17,12 @@ TeleFuser is a high-performance runtime for world model inference and multimodal ## News 📰 +- ✨ **2026-08-03**: Validated LingBot-World v2 target-side real-time generation on **4 x H100 80 GB** at + 832x480 and 16 FPS. The current 77-frame gate reached **17.14 steady compute FPS**; see the + [reproducible benchmark](docs/en/benchmark_aiperf.md#current-four-h100-real-time-gate). - ✨ **2026-07-27**: Unified streaming on LiveKit with room sessions, retained multi-session admission, LingBot chunk-boundary time slicing, reconnect-friendly browser transport, and server-push/bidirectional contracts. -- ✨ **2026-07-22**: **NEW** Added [**LingBot-Video**](examples/lingbot_video/README.md) support for Dense and MoE T2I/T2V/TI2V generation, native four-GPU CFG/SP execution, and in-memory MoE refinement. +- ✨ **2026-07-22**: Added [**LingBot-Video**](examples/lingbot_video/README.md) support for Dense and MoE T2I/T2V/TI2V generation, native four-GPU CFG/SP execution, and in-memory MoE refinement. - ✨ **2026-07-15**: Added [**LingBot-World v2**](https://github.com/Robbyant/lingbot-world-v2) support for offline generation, interactive WebRTC streaming, and multi-GPU inference. - ✨ **2026-07-06**: Added external **CacheSeek** latent cache integration for service-mode cross-request reuse. Cache hits can skip the first N denoising steps; the Wan2.2 cache-enabled service example snapshots `[5, 10, 15, 20, 25]` by default. See [docs/en/latent_cache.md](docs/en/latent_cache.md). @@ -95,6 +98,12 @@ video = pipe( TeleFuser streams `LingBot-World v2` through LiveKit. LingBot-World v2 uses camera control and its v2 PPL defaults; its streaming example caps a session at two minutes. +The validated four-H100 configuration sustains 17.14 target-side compute FPS for the default 77-frame, 832x480 +request, above its 16 FPS playback target. This is a synchronized pipeline-compute metric; model loading, LiveKit +encoding, network delivery, and client rendering are measured separately. See the +[LingBot example guide](examples/lingbot/README.md#validated-four-h100-real-time-gate) for the exact command and +chunk timings. + LingBot streaming uses the actor-based scheduler for both offline and service execution. Encode, DiT, and decode may overlap even on the same GPU; move stages only when memory placement requires it. See the [streaming scheduler guide](docs/en/stream_scheduler.md). diff --git a/README_zh.md b/README_zh.md index e88d415..09e238d 100644 --- a/README_zh.md +++ b/README_zh.md @@ -17,9 +17,12 @@ TeleFuser 是一个面向世界模型推理与多模态生成的高性能运行 ## News 📰 +- ✨ **2026-08-03**：LingBot-World v2 已在 **4 张 H100 80 GB** 上通过 832x480、16 FPS 的目标侧实时生成 + 验证。当前 77 帧门禁达到 **17.14 steady compute FPS**，复现方法见 + [基准文档](docs/zh/benchmark_aiperf.md#当前-77-帧实时计算门禁)。 - ✨ **2026-07-27**：统一使用 LiveKit 流式后端，支持 room 会话、worker 准入、浏览器自动重连，以及 server-push 和 bidirectional 两种 pipeline contract。 -- ✨ **2026-07-22**：**NEW** 新增 [**LingBot-Video**](examples/lingbot_video/README.md) 支持，覆盖 Dense/MoE T2I、T2V、TI2V、原生四卡 CFG/SP 推理与内存直传 MoE refiner。 +- ✨ **2026-07-22**：新增 [**LingBot-Video**](examples/lingbot_video/README.md) 支持，覆盖 Dense/MoE T2I、T2V、TI2V、原生四卡 CFG/SP 推理与内存直传 MoE refiner。 - ✨ **2026-07-15**：新增 [**LingBot-World v2**](https://github.com/Robbyant/lingbot-world-v2) 支持，支持离线生成、交互式 WebRTC 流和多卡推理。 - ✨ **2026-07-06**：新增外部 **CacheSeek** latent cache 集成，支持服务模式下跨请求复用；命中后可跳过前 N 步去噪。Wan2.2 服务示例默认快照 `[5, 10, 15, 20, 25]`。配置和安装方式见 [docs/zh/latent_cache.md](docs/zh/latent_cache.md)。 @@ -94,6 +97,10 @@ video = pipe( TeleFuser 通过 LiveKit 传输 `LingBot-World v2`。LingBot-World v2 使用相机控制和 v2 PPL 默认值；其流式 示例将单个会话上限设为两分钟。 +已验证的四卡 H100 配置在默认 77 帧、832x480 请求上达到 17.14 target-side compute FPS，高于 16 FPS +播放目标。该数值是设备同步后的 pipeline 计算指标；模型加载、LiveKit 编码、网络交付和客户端渲染需单独 +衡量。精确命令和逐 chunk 结果见 [LingBot 示例文档](examples/lingbot/README.md#validated-four-h100-real-time-gate)。 + LingBot 的离线与服务执行共用 actor scheduler。即使位于同一张 GPU，encode、DiT 和 decode 也可以重叠； 仅在显存放置需要时移动 Stage。详见[流式调度器指南](docs/zh/stream_scheduler.md)。 diff --git a/docs/en/benchmark_aiperf.md b/docs/en/benchmark_aiperf.md index f0cf212..6d24332 100644 --- a/docs/en/benchmark_aiperf.md +++ b/docs/en/benchmark_aiperf.md @@ -102,35 +102,59 @@ Target facts follow these rules: Client delivery, target pipeline residence, target phase time, and resource utilization remain separate dimensions. Fields without equivalent semantics remain private or unavailable instead of being forced into a common metric. -## Validated one-minute LingBot-World v2 replay - -Commit `663c385b179012c5c3de613212d10e8e6eac5f5d` was validated on 2026-08-02 with the -`stream_lingbot_world_v2_1min.json` workload, AIPerf 0.11.0 at commit -`e977ffbb1648510acec431b2a3fbd1a0f7bb8a35`, and four H100 80 GB GPUs. The current H100 example used BF16 DiT, -FP32 VAE, FlashAttention-4, disabled `torch.compile`, disabled FSDP, `chunk_size=4`, and 16 FPS. The 60-second -request was truncated to 60 complete latent chunks: 957 generated frames representing 59.75 seconds of media. -LingBot-World v2 used `local_attn_size=18` and `sink_size=6`; the session reported a fixed 28,080-token KV capacity -for its 240 latent frames. - -| Runtime / target | Compute FPS | Mean / p99 chunk | Stream FPS | Client frames | Artifact | -|---|---:|---:|---:|---:|---| -| TeleFuser `.venv`, torch cu128 | 16.191 | 0.988 / 1.099 s | 12.697 | 756 | `20260802_084922_d7ae0931` | -| TeleFuser `.venv-sglang`, torch cu130 | 15.897 | 1.006 / 1.208 s | 14.089 | 871 | `20260802_090301_af6c433c` | -| SGLang `.venv-sglang`, torch cu130 | 16.617 | 0.963 / 0.974 s | 16.772 | 957 | `20260801_104829_2320fd7f` | - -Every row completed 60 target chunks and generated 957 frames. AIPerf excluded only target chunk 0, leaving 944 -frames across 59 chunks. The aligned TeleFuser run used 59.381647 seconds of synchronized compute time and was 4.33% -below the SGLang compute rate. The cu130 TeleFuser result was 1.81% below its cu128 run, so the environment change is -reported separately and is not counted as an optimization gain. The aligned TeleFuser report is -`artifacts/telefuser_aiperf/stream_lingbot_v2_1min/20260802_090301_af6c433c/stream_report.html`. - -`stream_fps` is not used for the compute comparison. TeleFuser published LiveKit video with real-time 16 FPS pacing; -its aligned run averaged 18.99 ms from decoded-ready to publish start, 941.66 ms in paced publication, and 2.10 ms -from publish completion to client metadata. SGLang used unpaced burst WebSocket output. Those delivery semantics are -not equivalent even though both include network and client decoding. TeleFuser's 9.740-second first-frame latency -comprised 0.630 seconds to create the session, another 1.979 seconds to connect, 3.206 seconds from connection to -admission, and 3.925 seconds from admission to the first client frame; runtime creation occupied 1.564 seconds of the -last interval. +## Four-H100 LingBot-World v2 validation + +Both runs below used four H100 80 GB GPUs, BF16 DiT, FP32 VAE, FlashAttention-4, disabled FSDP, disabled +`torch.compile`, `chunk_size=4`, and 16 FPS output. They validate different workloads and code revisions, so their +results should not be interpreted as a before-and-after performance comparison. + +### Current four-H100 real-time gate + +Commit `540b579` was validated on 2026-08-03 through the direct LingBot pipeline-service path with four H100 80 GB +GPUs and PyTorch 2.11.0+cu128. The request used 832x480 output, 77 frames, and five four-latent-frame chunks. + +| Metric | Result | +|---|---:| +| Generated frames / target chunks | 77 / 5 | +| Steady chunks after excluding chunk 0 | 4 | +| Steady compute FPS | **17.1399** | +| Chunk compute mean / p50 / p90 / max | 0.9335 / 0.9409 / 0.9410 / 1.0058 s | +| First generated frame from measured session start | 3.2182 s | + +This run clears the average target-side 16 FPS compute gate. It does not claim that every chunk clears the one-second +budget: the maximum was 1.0058 seconds. The synchronized compute interval includes condition handling, DiT, +clean-KV update, spatial VAE decode, GPU-to-CPU transfer, and frame conversion. It excludes model loading, runtime +creation, LiveKit pacing/encoding, network delivery, and client rendering. + +The exact reproduction command is in the +[LingBot example guide](https://github.com/Tele-AI/TeleFuser/tree/main/examples/lingbot#validated-four-h100-real-time-gate) +and uses `tools/validation/benchmark_lingbot_world_v2_direct.py`. + +### Current one-minute streaming replay + +The one-minute workload was rerun on 2026-08-03 at TeleFuser commit +`284996dd616cfd44a55523687b7f2a63a281abb9`. It validates sustained target generation, bounded KV-cache capacity, +and the paced LiveKit delivery path on the current communication-optimized revision. + +The run used the `stream_lingbot_world_v2_1min.json` workload and AIPerf 0.11.0 at commit +`e977ffbb1648510acec431b2a3fbd1a0f7bb8a35`. The 60-second request was truncated to 60 complete latent chunks: +957 generated frames representing 59.75 seconds of media. With `local_attn_size=18` and `sink_size=6`, the +240-latent-frame session reported a fixed 28,080-token KV capacity. + +| Metric | Result | +|---|---:| +| Generated target frames / chunks | 957 / 60 | +| Steady frames / chunks after excluding chunk 0 | 944 / 59 | +| Steady target compute time / FPS | 58.2791 s / **16.1979** | +| Chunk compute mean / p50 / p90 / p99 / max | 0.9878 / 0.9593 / 1.0624 / 1.0932 / 1.1149 s | +| LiveKit stream FPS / client frames | 13.1967 / 803 | +| First client frame / session runtime | 6.0682 / 66.8948 s | +| Runtime creation | 1.4176 s | +| Artifact | `20260803_095518_62ec043c` | + +The target completed all 60 chunks and cleared the average 16 FPS compute gate. It did not keep every chunk below one +second: p99 was 1.0932 seconds and the maximum was 1.1149 seconds. The lower client frame count belongs to the paced +delivery measurement and must not be conflated with target generation completeness. ## Reproducibility diff --git a/docs/zh/benchmark_aiperf.md b/docs/zh/benchmark_aiperf.md index 26e1c65..45da696 100644 --- a/docs/zh/benchmark_aiperf.md +++ b/docs/zh/benchmark_aiperf.md @@ -97,31 +97,55 @@ Target 原始事实遵守以下规则： 客户端交付、target pipeline residence、target phase time 和资源利用率保持为不同维度。无法等价的字段保留为 private 或 unavailable，不强行映射为同一指标。 -## LingBot-World v2 一分钟回放实测 - -2026-08-02 使用 4 张 H100 80 GB 验证了 TeleFuser commit -`663c385b179012c5c3de613212d10e8e6eac5f5d` 和 `stream_lingbot_world_v2_1min.json` workload；AIPerf 为 -0.11.0、commit `e977ffbb1648510acec431b2a3fbd1a0f7bb8a35`。当前 H100 example 使用 BF16 DiT、FP32 VAE、 -FlashAttention-4，关闭 `torch.compile` 和 FSDP，`chunk_size=4`，输出 16 FPS。60 秒请求按完整 latent -chunk 截断为 60 个 chunk、957 帧，对应 59.75 秒媒体时长。LingBot-World v2 使用 -`local_attn_size=18`、`sink_size=6`，本次 240 latent frame session 报告的固定 KV 容量为 28,080 token。 - -| 运行环境 / target | Compute FPS | Chunk mean / p99 | Stream FPS | 客户端帧数 | Artifact | -|---|---:|---:|---:|---:|---| -| TeleFuser `.venv`，torch cu128 | 16.191 | 0.988 / 1.099 s | 12.697 | 756 | `20260802_084922_d7ae0931` | -| TeleFuser `.venv-sglang`，torch cu130 | 15.897 | 1.006 / 1.208 s | 14.089 | 871 | `20260802_090301_af6c433c` | -| SGLang `.venv-sglang`，torch cu130 | 16.617 | 0.963 / 0.974 s | 16.772 | 957 | `20260801_104829_2320fd7f` | - -三次运行均完成 60 个 target chunk、生成 957 帧。AIPerf 只排除 target chunk 0，稳态统计包含 59 个 chunk、 -944 帧。对齐环境后的 TeleFuser 同步计算时间为 59.381647 秒，compute FPS 比 SGLang 低 4.33%。TeleFuser -cu130 结果比 cu128 低 1.81%，因此环境变化单独报告，不计作代码优化收益。对齐环境 TeleFuser 报告位于 -`artifacts/telefuser_aiperf/stream_lingbot_v2_1min/20260802_090301_af6c433c/stream_report.html`。 - -Compute 对比不使用 `stream_fps`。TeleFuser 的 LiveKit 视频按实时 16 FPS pacing 发布；对齐环境运行中， -decoded-ready 到 publish start 平均 18.99 ms，paced publish 平均 941.66 ms，publish 完成到客户端 metadata -平均 2.10 ms。SGLang 使用无 pacing 的 WebSocket burst 输出，两者交付语义不等价，尽管两边都包含网络 -传输和客户端解码。TeleFuser 首帧为 9.740 秒：session 创建 0.630 秒，其后连接 1.979 秒，连接到准入 -3.206 秒，准入到客户端首帧 3.925 秒；最后一段中的 runtime creation 为 1.564 秒。 +## 四卡 H100 LingBot-World v2 验证 + +以下两次运行均使用 4 张 H100 80 GB、BF16 DiT、FP32 VAE、FlashAttention-4，关闭 FSDP 和 +`torch.compile`，设置 `chunk_size=4` 并输出 16 FPS。两次运行的 workload 和代码版本不同，不能将结果 +解释为优化前后的性能对比。 + +### 当前 77 帧实时计算门禁 + +2026-08-03 使用 PyTorch 2.11.0+cu128，通过 direct LingBot pipeline-service 路径验证了 commit +`540b579`。请求使用 832x480、77 帧，共生成 5 个、每个包含 4 个 latent frame 的 chunk。 + +| 指标 | 结果 | +|---|---:| +| 生成帧数 / target chunk | 77 / 5 | +| 排除 chunk 0 后的 steady chunk | 4 | +| Steady compute FPS | **17.1399** | +| Chunk compute mean / p50 / p90 / max | 0.9335 / 0.9409 / 0.9410 / 1.0058 秒 | +| 从计时 session 开始到首个生成帧 | 3.2182 秒 | + +该运行通过了平均 target-side 16 FPS 计算门禁，但不表示每个 chunk 都低于一秒：最大值为 1.0058 秒。 +设备同步后的 compute 区间包含 condition handling、DiT、clean-KV update、空间 VAE decode、GPU-to-CPU +传输和 frame conversion；不包含模型加载、runtime creation、LiveKit pacing/encoding、网络交付和客户端渲染。 + +复现命令见 [LingBot 示例文档](https://github.com/Tele-AI/TeleFuser/tree/main/examples/lingbot#validated-four-h100-real-time-gate)， +使用 `tools/validation/benchmark_lingbot_world_v2_direct.py`。 + +### 当前一分钟流式回放 + +2026-08-03 在 TeleFuser commit `284996dd616cfd44a55523687b7f2a63a281abb9` 上重新运行了一分钟 workload， +用于验证当前通信优化版本的持续 target 生成、固定 KV cache 容量和带 pacing 的 LiveKit 交付路径。 + +运行使用 `stream_lingbot_world_v2_1min.json` workload，以及 commit +`e977ffbb1648510acec431b2a3fbd1a0f7bb8a35` 对应的 AIPerf 0.11.0。60 秒请求按完整 latent chunk 截断为 +60 个 chunk、957 帧，对应 59.75 秒媒体时长。使用 `local_attn_size=18` 和 `sink_size=6` 时，240 个 +latent frame 的 session 报告固定 KV 容量为 28,080 token。 + +| 指标 | 结果 | +|---|---:| +| Target 生成帧数 / chunk | 957 / 60 | +| 排除 chunk 0 后的 steady 帧数 / chunk | 944 / 59 | +| Steady target compute 时间 / FPS | 58.2791 秒 / **16.1979** | +| Chunk compute mean / p50 / p90 / p99 / max | 0.9878 / 0.9593 / 1.0624 / 1.0932 / 1.1149 秒 | +| LiveKit stream FPS / 客户端帧数 | 13.1967 / 803 | +| 客户端首帧 / session runtime | 6.0682 / 66.8948 秒 | +| Runtime creation | 1.4176 秒 | +| Artifact | `20260803_095518_62ec043c` | + +Target 完成全部 60 个 chunk，平均 compute 通过 16 FPS 门禁，但并非每个 chunk 都低于一秒：p99 为 +1.0932 秒，最大值为 1.1149 秒。客户端帧数较低属于带 pacing 的交付测量，不能与 target 生成完整性混为一谈。 ## 复现要求 diff --git a/examples/lingbot/README.md b/examples/lingbot/README.md index 30968d2..f7e515d 100644 --- a/examples/lingbot/README.md +++ b/examples/lingbot/README.md @@ -125,6 +125,43 @@ python examples/lingbot/lingbot_world_v2_image_to_video_h100.py \ --v2_model_root "${TF_MODEL_ZOO_PATH}/lingbot/lingbot-world-v2-14b-causal-fast/transformers" ``` +### Validated Four-H100 Real-Time Gate + +Commit `540b579` was validated on 2026-08-03 with four H100 80 GB GPUs, PyTorch 2.11.0+cu128, +FlashAttention-4, BF16 DiT, FP32 VAE, disabled FSDP, and disabled `torch.compile`. The default 832x480 +request generated all 77 frames in five chunks at a 16 FPS playback target. + +| Metric | Result | +| --- | ---: | +| Steady compute FPS | **17.14** | +| Steady chunk mean / p50 / p90 | 0.9335 / 0.9409 / 0.9410 s | +| Slowest steady chunk | 1.0058 s | +| Generated frames / chunks | 77 / 5 | + +The steady summary excludes chunk 0 and covers four 16-frame chunks. `compute_seconds` synchronizes all target CUDA +devices and includes condition handling, DiT, clean-KV update, spatial VAE decode, GPU-to-CPU transfer, and frame +conversion. It excludes model loading, runtime creation, LiveKit pacing/encoding, network delivery, and client +rendering. The average therefore clears the 16 FPS target-side real-time gate, while the slowest chunk exceeds its +one-second budget by 5.8 ms; treat this as a validated configuration, not a guarantee for other hardware, resolutions, +durations, concurrent sessions, or transport conditions. + +Reproduce the measured direct pipeline-service path without LiveKit or codec time: + +```bash +TF_MODEL_ZOO_PATH=/path/to/model_zoo \ +CUDA_VISIBLE_DEVICES=0,1,2,3 \ +python tools/validation/benchmark_lingbot_world_v2_direct.py \ + --pipeline examples/lingbot/lingbot_world_v2_image_to_video_h100.py \ + --image examples/data/lingbot_world_fast/image.jpg \ + --control-trace benchmarks/telefuser_aiperf/data/stream_lingbot_controls.json \ + --output work_dirs/lingbot_world_v2_4gpu_77frames.json \ + --gpu-num 4 --frame-num 77 --fps 16 --chunk-size 4 +``` + +The offline CLI was also validated to produce an H.264 832x480 video containing all 77 frames. Use the +[AIPerf benchmark guide](../../docs/en/benchmark_aiperf.md) for the one-minute workload, client delivery metrics, +and comparisons that require identical environments. + ## Usage ### Four H100 GPUs