From 9005d789033b8c3ec5876a7a68c4e2d9238f5c69 Mon Sep 17 00:00:00 2001 From: linzhenpl07 <60651557+linzhenpl07@users.noreply.github.com> Date: Wed, 16 Sep 2026 15:01:11 +0800 Subject: [PATCH] [LingBot World] Reuse the text encode for a prompt already encoded (#6845) Signed-off-by: linzhenpl07 Co-authored-by: Claude Opus 5 --- .../test_pipeline_lingbot_world.py | 122 ++++++++++++++++++ .../models/lingbot_world/pipeline.py | 34 ++++- 2 files changed, 154 insertions(+), 2 deletions(-) diff --git a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py index fb9e93a5d..9d5e958c5 100644 --- a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py +++ b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py @@ -2711,3 +2711,125 @@ def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> No assert "from .pipeline import" in lingbot_init assert '"LingBotWorldCausalDMDPipeline"' in lingbot_init assert '"CausalLingBotWorldTransformer3DModel"' in lingbot_init + + +# --------------------------------------------------------------------------- +# a prompt is encoded once, whichever path asks for it +# --------------------------------------------------------------------------- +def _counting_text_encoder(pipeline) -> list[str]: + """Give ``pipeline`` its real encode_prompt over stubs that record each encode.""" + del pipeline.encode_prompt # _pipeline() replaces it; these tests need the real one + encoded: list[str] = [] + + def tokenizer(texts, **kwargs): + del kwargs + encoded.extend(texts) + return SimpleNamespace( + input_ids=torch.arange(512).view(1, 512), + attention_mask=torch.ones(1, 512, dtype=torch.long), + ) + + def text_encoder(input_ids, attention_mask): + del input_ids, attention_mask + return SimpleNamespace(last_hidden_state=torch.full((1, 512, 8), float(len(encoded)))) + + pipeline.tokenizer = tokenizer + pipeline.text_encoder = text_encoder + return encoded + + +def test_the_text_encoder_runs_once_per_distinct_prompt() -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + encoded = _counting_text_encoder(pipeline) + + def encode(prompt, *, length=512, dtype=torch.float32): + return pipeline.encode_prompt(prompt, max_sequence_length=length, dtype=dtype) + + first = encode("move through the room") + assert encode("move through the room") is first + # The key is the text the tokenizer actually sees, so spacing alone is not new. + assert encode(" move through the room ") is first + assert encoded == ["move through the room"] + + encode("enter the snowy valley") + # dtype and sequence length change the tensor, so each is its own entry. + encode("move through the room", dtype=torch.float16) + encode("move through the room", length=256) + assert encoded == [ + "move through the room", + "enter the snowy valley", + "move through the room", + "move through the room", + ] + + +def test_the_prompt_encode_cache_is_bounded() -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + encoded = _counting_text_encoder(pipeline) + size = module._PROMPT_EMBEDS_CACHE_SIZE + + for index in range(size + 1): + pipeline.encode_prompt(f"prompt {index}", max_sequence_length=512, dtype=torch.float32) + assert len(pipeline._prompt_embeds_cache) == size + + # The oldest entry was the one evicted, so it is the one encoded again. + pipeline.encode_prompt("prompt 0", max_sequence_length=512, dtype=torch.float32) + pipeline.encode_prompt(f"prompt {size}", max_sequence_length=512, dtype=torch.float32) + assert encoded == [f"prompt {index}" for index in range(size + 1)] + ["prompt 0"] + + +def test_realtime_ticks_reuse_the_encode_and_still_invalidate_on_a_prompt_change( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + pipeline._ar_height = 16 + pipeline._ar_width = 16 + pipeline._ar_diffusion_kv_state = object() + encoded = _counting_text_encoder(pipeline) + invalidations: list[bool] = [] + + def ar_text_caches(prompt_embeds, *, invalidate): + del prompt_embeds + invalidations.append(invalidate) + return [SimpleNamespace()] + + monkeypatch.setattr(pipeline, "_ar_text_caches", ar_text_caches) + monkeypatch.setattr( + pipeline, + "_generate_block", + lambda **kwargs: torch.randn((1, 16, 3, 2, 2), generator=kwargs["generator"]), + ) + + for chunk_index in range(3): + pipeline(_request(sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=chunk_index)))) + switched = _prompt() + switched["prompt"] = "enter the snowy valley" + pipeline( + _request( + sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=3, prompt=switched["prompt"])), + prompt=switched, + ) + ) + + assert encoded == ["move through the room", "enter the snowy valley"] + # A new prompt must still drop the cross-attention K/V built from the old one. + assert invalidations == [False, False, False, True] + + +def test_stepwise_requests_with_the_same_prompt_share_one_encode() -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module, transformer=_RecordingTransformer()) + pipeline._ar_height = 16 + pipeline._ar_width = 16 + encoded = _counting_text_encoder(pipeline) + + states = [_stepwise_state(request_id=request_id, num_frames=21) for request_id in ("req-a", "req-b")] + for state in states: + with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): + pipeline.prepare_encode(state) + + assert encoded == ["move through the room"] + assert states[0].prompt_embeds is states[1].prompt_embeds diff --git a/vllm_omni/diffusion/models/lingbot_world/pipeline.py b/vllm_omni/diffusion/models/lingbot_world/pipeline.py index ccaa677e0..8e8d5dea4 100644 --- a/vllm_omni/diffusion/models/lingbot_world/pipeline.py +++ b/vllm_omni/diffusion/models/lingbot_world/pipeline.py @@ -6,6 +6,7 @@ from __future__ import annotations import math import os +from collections import OrderedDict from collections.abc import Callable, Iterable, Iterator, Sequence from contextlib import contextmanager, nullcontext from dataclasses import dataclass, replace @@ -99,6 +100,10 @@ _CAMERA_SPATIAL_FOLD = 8 _MAX_PIXEL_AREA = 480 * 832 _MAX_SOURCE_IMAGE_PIXELS = 4096 * 4096 _MAX_SEQUENCE_LENGTH = 512 +# Distinct prompts whose text encodes are kept. One entry is a single padded +# sequence -- a few MiB at 512 tokens -- so a handful covers a session's prompt +# plus the switches it makes, at no meaningful cost next to its KV. +_PROMPT_EMBEDS_CACHE_SIZE = 4 _ACTION_ROOT_ENV = "VLLM_OMNI_LINGBOT_ACTION_ROOT" _PREPROCESSED_CAMERA_KEY = "_lingbot_camera_trajectory" _PREPROCESSED_CAMERA_ACTIONS_KEY = "_lingbot_camera_actions" @@ -569,6 +574,8 @@ class LingBotWorldCausalDMDPipeline( self._ar_width = int(model_config.get("ar_diffusion_width", 832)) self._ar_diffusion_kv_state: ARDiffusionKVState | None = None self._ar_sessions: dict[str, _LingBotARSessionState] = {} + # Text encodes keyed by what they depend on; see encode_prompt. + self._prompt_embeds_cache: OrderedDict[tuple[str, int, torch.dtype], torch.Tensor] = OrderedDict() # One temporal decoder cache per stepwise session, keyed the way the # runner keys AR sessions (session_id == request_id) so both are # released together rather than through two independent lifecycles. @@ -1221,8 +1228,27 @@ class LingBotWorldCausalDMDPipeline( max_sequence_length: int, dtype: torch.dtype, ) -> torch.Tensor: + """Encode ``prompt`` with UMT5, reusing the result for a prompt already encoded. + + The encode depends only on the whitespace-normalised text, the sequence + length and the dtype -- the tokenizer and text encoder are fixed once + loaded -- so those three are the key. That covers every caller the same + way: a realtime tick repeats its session's prompt on every block, a + stepwise request encodes once per request, and a server tends to open + many sessions with the same scene prompt. + + A hit returns the stored tensor itself. Callers only read it: the + cross-attention projection and the DMD transformer both work out of + place, so a shared encode is never changed under another session. + """ + text = " ".join(prompt.strip().split()) + key = (text, int(max_sequence_length), dtype) + cached = self._prompt_embeds_cache.get(key) + if cached is not None: + self._prompt_embeds_cache.move_to_end(key) + return cached text_inputs = self.tokenizer( - [" ".join(prompt.strip().split())], + [text], padding="max_length", max_length=max_sequence_length, truncation=True, @@ -1234,7 +1260,11 @@ class LingBotWorldCausalDMDPipeline( attention_mask = text_inputs.attention_mask.to(self.device) prompt_embeds = self.text_encoder(input_ids, attention_mask).last_hidden_state prompt_embeds = prompt_embeds.to(device=self.device, dtype=dtype) - return prompt_embeds * attention_mask.unsqueeze(-1).to(dtype=prompt_embeds.dtype) + prompt_embeds = prompt_embeds * attention_mask.unsqueeze(-1).to(dtype=prompt_embeds.dtype) + self._prompt_embeds_cache[key] = prompt_embeds + if len(self._prompt_embeds_cache) > _PROMPT_EMBEDS_CACHE_SIZE: + self._prompt_embeds_cache.popitem(last=False) + return prompt_embeds def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: inputs = self._parse_request(req)