From 8efd3c2f6dbf3fd3852168d77852c3528c7932fc Mon Sep 17 00:00:00 2001 From: wtz2333 <2955110911@qq.com> Date: Mon, 21 Sep 2026 17:02:15 +0800 Subject: [PATCH] [Feature] Add experimental LingBot World last-step KV reuse (#7816) Signed-off-by: wtz2333 <2955110911@qq.com> --- recipes/Robbyant/LingBot-World-2.0.md | 20 +++++++ .../test_pipeline_lingbot_world.py | 54 +++++++++++++++---- .../models/lingbot_world/dmd_block.py | 35 ++++++++---- .../models/lingbot_world/pipeline.py | 7 ++- 4 files changed, 94 insertions(+), 22 deletions(-) diff --git a/recipes/Robbyant/LingBot-World-2.0.md b/recipes/Robbyant/LingBot-World-2.0.md index 24a36653d..7add2f7d8 100644 --- a/recipes/Robbyant/LingBot-World-2.0.md +++ b/recipes/Robbyant/LingBot-World-2.0.md @@ -53,6 +53,26 @@ source frames, at least as many as requested and at most 4096. The runtime consumes only the prefix needed by the request. Longer rollouts also require sufficient device memory. +## Experimental KV writeback reuse + +To reuse KV from the final denoising step instead of running the additional +clean-latent writeback, pass this option in the `model_config` supplied to +`Omni` or `AsyncOmni`: + +```python +model_config = {"lingbot_reuse_last_step_kv": True} +``` + +The default is `False`. The value must be a boolean, not a string or integer. +It is validated and fixed when the pipeline instance is constructed; changing +it requires a new instance. It is not a per-request `extra_args` option. +The previous `VLLM_OMNI_LINGBOT_REUSE_LAST_STEP_KV` environment variable is no +longer read. + +This experimental mode stores the final noisy probe's KV rather than clean +`x0` KV, so generated outputs can differ. The four denoising steps, chunk +boundaries, and paged-context finalization remain the same. + ## Realtime in-process generation (deprecated) Prefer [Streaming video serving](#streaming-video-serving) for mid-session 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 91ae00b7d..1ccc7aadc 100644 --- a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py +++ b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py @@ -1358,10 +1358,34 @@ def test_first_frame_condition_and_camera_fold_match_transformer_contract() -> N assert expected_camera.shape == (1, 384, 3, 2, 2) -def test_fixed_dmd_transition_and_cache_commit_trace() -> None: +@pytest.mark.parametrize("value", ["true", "false", 1, 0, None]) +def test_reuse_last_step_kv_rejects_non_boolean_config(value: object) -> None: + config = _od_config() + config.model_config["lingbot_reuse_last_step_kv"] = value + with pytest.raises(ValueError, match="lingbot_reuse_last_step_kv must be a bool"): + _pipeline(_load_pipeline_module(), od_config=config) + + +def test_reuse_last_step_kv_is_fixed_per_instance(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("VLLM_OMNI_LINGBOT_REUSE_LAST_STEP_KV", "1") + module = _load_pipeline_module() + config = _od_config() + config.model_config["lingbot_reuse_last_step_kv"] = True + reuse = _pipeline(module, od_config=config) + default = _pipeline(module) + config.model_config["lingbot_reuse_last_step_kv"] = False + assert reuse._dmd_blocks.reuse_last_step_kv is True + assert default._dmd_blocks.reuse_last_step_kv is False + assert reuse._dmd_blocks is reuse._dmd_blocks + + +@pytest.mark.parametrize("reuse_last_step_kv", [False, True]) +def test_fixed_dmd_transition_and_cache_commit_trace(reuse_last_step_kv: bool) -> None: module = _load_pipeline_module() transformer = _RecordingTransformer() - pipeline = _pipeline(module, transformer=transformer) + config = _od_config() + config.model_config["lingbot_reuse_last_step_kv"] = reuse_last_step_kv + pipeline = _pipeline(module, transformer=transformer, od_config=config) result = pipeline(_request()) @@ -1383,14 +1407,18 @@ def test_fixed_dmd_transition_and_cache_commit_trace() -> None: current = x0 torch.testing.assert_close(result.output, current) + expected_timesteps = [timestep for timestep, _ in warped_schedule] + if not reuse_last_step_kv: + expected_timesteps.append(0.0) torch.testing.assert_close( - torch.cat([call["timestep"] for call in transformer.calls]), - torch.tensor([*(timestep for timestep, _ in warped_schedule), 0.0]), + torch.cat([call["timestep"] for call in transformer.calls]), torch.tensor(expected_timesteps) ) - assert [call["update_cache"] for call in transformer.calls] == [False, False, False, False, True] - assert [call["start_frame"] for call in transformer.calls] == [0, 0, 0, 0, 0] + commit_flags = [False] * (len(expected_timesteps) - 1) + [True] + assert [call["update_cache"] for call in transformer.calls] == commit_flags + assert [call["start_frame"] for call in transformer.calls] == [0] * len(expected_timesteps) assert len({call["cache_id"] for call in transformer.calls}) == 1 - torch.testing.assert_close(transformer.calls[-1]["hidden_states"][:, :16], result.output) + committed_latent = result.output + warped_schedule[-1][1] if reuse_last_step_kv else result.output + torch.testing.assert_close(transformer.calls[-1]["hidden_states"][:, :16], committed_latent) @pytest.mark.parametrize( @@ -2389,10 +2417,13 @@ def test_forward_rejects_a_stepwise_camera_action_script(extra_args) -> None: pipeline(_request(sampling=sampling)) -def test_stepwise_progress_metadata_and_commit_trace() -> None: +@pytest.mark.parametrize("reuse_last_step_kv", [False, True]) +def test_stepwise_progress_metadata_and_commit_trace(reuse_last_step_kv: bool) -> None: module = _load_pipeline_module() transformer = _RecordingTransformer() - pipeline = _pipeline(module, transformer=transformer) + config = _od_config() + config.model_config["lingbot_reuse_last_step_kv"] = reuse_last_step_kv + pipeline = _pipeline(module, transformer=transformer, od_config=config) pipeline._ar_height = 16 pipeline._ar_width = 16 state = _stepwise_state(num_frames=21) @@ -2406,8 +2437,9 @@ def test_stepwise_progress_metadata_and_commit_trace() -> None: assert [output.chunk_index for output in outputs] == [0, 1] assert all(output.total_chunks == 2 for output in outputs) assert outputs[-1].finished is True - assert [call["update_cache"] for call in transformer.calls] == [False, False, False, False, True] * 2 - assert [call["start_frame"] for call in transformer.calls] == [0] * 5 + [3] * 5 + calls_per_chunk = 4 if reuse_last_step_kv else 5 + assert [call["update_cache"] for call in transformer.calls] == ([False] * (calls_per_chunk - 1) + [True]) * 2 + assert [call["start_frame"] for call in transformer.calls] == [0] * calls_per_chunk + [3] * calls_per_chunk assert fake.commits == ["main", "main"] for chunk_index, output in enumerate(outputs): metadata = output.output["metadata"]["ar_diffusion"] diff --git a/vllm_omni/diffusion/models/lingbot_world/dmd_block.py b/vllm_omni/diffusion/models/lingbot_world/dmd_block.py index b179eb773..27da62149 100644 --- a/vllm_omni/diffusion/models/lingbot_world/dmd_block.py +++ b/vllm_omni/diffusion/models/lingbot_world/dmd_block.py @@ -7,9 +7,10 @@ over every block, stepwise execution calls ``probe_step`` / ``apply_transition`` one denoise step at a time and ``commit_block_kv`` from ``post_decode`` -- so the math lives here once and cannot drift between the two modes. -A block is four probes that must not touch KV, then one clean-x0 commit that -does. Whether a call sees paged (session-bound) or request-local KV is decided -by the ``ar`` argument, so this class never reaches back into the pipeline. +By default a block is four probes followed by one clean-x0 KV commit. +The experimental last-step reuse mode instead commits the fourth noisy probe. +Whether a call sees paged (session-bound) or request-local KV is decided by +the ``ar`` argument, so this class never reaches back into the pipeline. """ from __future__ import annotations @@ -19,6 +20,7 @@ from typing import TYPE_CHECKING, Any, cast import torch from diffusers.utils.torch_utils import randn_tensor +from vllm.logger import init_logger from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx from vllm_omni.diffusion.models.lingbot_world.transformer import ( @@ -33,6 +35,9 @@ if TYPE_CHECKING: from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState +logger = init_logger(__name__) + + @dataclass(frozen=True) class ARBlockContext: """The bound AR-Diffusion session one block runs against. @@ -55,10 +60,14 @@ class LingBotDMDBlockRunner: *, device: torch.device, enforce_eager: bool, + reuse_last_step_kv: bool = False, ) -> None: self.transformer = transformer self.device = device self.enforce_eager = enforce_eager + # Experimental approximation: cache x_t/t_last features instead of x0/t=0. + self.reuse_last_step_kv = reuse_last_step_kv + logger.info("LingBot experimental last-step KV reuse: %s", self.reuse_last_step_kv) # ── cache selection ────────────────────────────────────────────────── @@ -114,8 +123,8 @@ class LingBotDMDBlockRunner: ) -> torch.Tensor: """Predict flow for one denoise step. - A probe never writes KV: only the clean x0 of a finished block may - enter the cache, which is ``commit_block_kv``'s job. + Normally only clean x0 enters KV. The experimental reuse mode commits + the final noisy probe instead; the fixed DMD schedule has four probes. ``camera_cache`` is the block's camera-modulation cache when the caller holds it across separate calls (the stepwise path); ``generate_block`` @@ -128,14 +137,15 @@ class LingBotDMDBlockRunner: # Checkpoint channel contract: # [noise/x_t(16), temporal_mask(4), image_latent(16)] -> 36. model_input = torch.cat((current_latents.to(dtype=condition.dtype), condition), dim=1) + commit_probe = self.reuse_last_step_kv and step_index == 3 flow_prediction = self.transformer( hidden_states=model_input, timestep=timestep, encoder_hidden_states=prompt_embeds, camera_hidden_states=camera, - cache=self._block_cache(condition=condition, cache=cache, ar=ar, commit_current=False), + cache=self._block_cache(condition=condition, cache=cache, ar=ar, commit_current=commit_probe), start_frame=start_frame, - update_cache=False, + update_cache=commit_probe, camera_modulation_cache=camera_cache, ) if flow_prediction.shape != current_latents.shape: @@ -183,10 +193,15 @@ class LingBotDMDBlockRunner: ) -> None: """Write the finished block's clean x0 into KV and commit its pages. - The fifth transformer call of a block, deliberately not a denoise step: - on the stepwise path it belongs to ``post_decode()``. ``camera_cache`` - is as in ``probe_step``. + By default this is the fifth transformer call, at t=0. Experimental + reuse only finalizes the pages written by the fourth probe. On the + stepwise path both modes finalize from ``post_decode()``. + ``camera_cache`` is as in ``probe_step``. """ + if self.reuse_last_step_kv: + if ar is not None: + ar.state.commit_paged_context(ar.branch) + return # Commit K/V only for the final clean block, never for noisy probes. cache_input = torch.cat((latents.to(dtype=condition.dtype), condition), dim=1) if not self.enforce_eager: diff --git a/vllm_omni/diffusion/models/lingbot_world/pipeline.py b/vllm_omni/diffusion/models/lingbot_world/pipeline.py index 48c2ca568..ae24c0138 100644 --- a/vllm_omni/diffusion/models/lingbot_world/pipeline.py +++ b/vllm_omni/diffusion/models/lingbot_world/pipeline.py @@ -581,6 +581,11 @@ class LingBotWorldCausalDMDPipeline( del prefix _validate_parallel_config(od_config) self.od_config = od_config + model_config = getattr(od_config, "model_config", None) or {} + reuse_last_step_kv = model_config.get("lingbot_reuse_last_step_kv", False) + if not isinstance(reuse_last_step_kv, bool): + raise ValueError("model_config.lingbot_reuse_last_step_kv must be a bool.") + self._reuse_last_step_kv = reuse_last_step_kv self.device = get_local_device() dtype = getattr(od_config, "dtype", torch.bfloat16) model = od_config.model @@ -666,7 +671,6 @@ class LingBotWorldCausalDMDPipeline( self.vae_scale_factor_temporal = int(getattr(self.vae.config, "scale_factor_temporal", 4)) self.vae_scale_factor_spatial = int(getattr(self.vae.config, "scale_factor_spatial", 8)) - model_config = getattr(od_config, "model_config", None) or {} self._ar_height = int(model_config.get("ar_diffusion_height", 480)) self._ar_width = int(model_config.get("ar_diffusion_width", 832)) self._ar_diffusion_kv_state: ARDiffusionKVState | None = None @@ -1337,6 +1341,7 @@ class LingBotWorldCausalDMDPipeline( self.transformer, device=self.device, enforce_eager=bool(self.od_config.enforce_eager), + reuse_last_step_kv=self._reuse_last_step_kv, ) self._dmd_block_runner = runner return runner