[Feature] Add experimental LingBot World last-step KV reuse (#7816)

Signed-off-by: wtz2333 <2955110911@qq.com>

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

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"]

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:

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