[LingBot World] Eleven bit-identical removals of repeated work in the realtime path (#7648)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

vllm_omni/diffusion/models/interface.py

@@ -92,6 +92,13 @@ class SupportsStepExecution(Protocol):
     ``prepare_encode()`` (one-time request setup), ``denoise_step()``
     (one denoise forward), ``step_scheduler()`` (one scheduler update),
     and ``post_decode()`` (final decode).
+
+    A pipeline may additionally set the optional class attribute
+    ``supports_chunk_step_grouping = True`` to declare that its request state
+    may advance through every denoise step of its current chunk without a
+    serving-scheduler cycle in between (a capability, not a policy: the runner
+    that owns the scheduling decision chooses whether to group steps). It is
+    optional so it does not become part of the runtime protocol check.
     """
 
     supports_step_execution: ClassVar[bool] = True

vllm_omni/diffusion/models/lingbot_world/dmd_block.py

@@ -22,6 +22,7 @@ from diffusers.utils.torch_utils import randn_tensor
 
 from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
 from vllm_omni.diffusion.models.lingbot_world.transformer import (
+    CameraModulationCache,
     LingBotAttentionCache,
     LingBotTransformerCache,
 )
@@ -109,11 +110,16 @@ class LingBotDMDBlockRunner:
         start_frame: int,
         timestep_value: float,
         step_index: int,
+        camera_cache: CameraModulationCache | None = None,
     ) -> 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.
+
+        ``camera_cache`` is the block's camera-modulation cache when the caller
+        holds it across separate calls (the stepwise path); ``generate_block``
+        opens the window itself instead.
         """
         if not self.enforce_eager:
             torch.compiler.cudagraph_mark_step_begin()
@@ -130,6 +136,7 @@ class LingBotDMDBlockRunner:
             cache=self._block_cache(condition=condition, cache=cache, ar=ar, commit_current=False),
             start_frame=start_frame,
             update_cache=False,
+            camera_modulation_cache=camera_cache,
         )
         if flow_prediction.shape != current_latents.shape:
             raise RuntimeError(
@@ -172,11 +179,13 @@ class LingBotDMDBlockRunner:
         cache: LingBotTransformerCache | None,
         ar: ARBlockContext | None,
         start_frame: int,
+        camera_cache: CameraModulationCache | None = None,
     ) -> 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()``.
+        on the stepwise path it belongs to ``post_decode()``. ``camera_cache``
+        is as in ``probe_step``.
         """
         # 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)
@@ -190,6 +199,7 @@ class LingBotDMDBlockRunner:
             cache=self._block_cache(condition=condition, cache=cache, ar=ar, commit_current=True),
             start_frame=start_frame,
             update_cache=True,
+            camera_modulation_cache=camera_cache,
         )
         if ar is not None:
             ar.state.commit_paged_context(ar.branch)
@@ -210,6 +220,10 @@ class LingBotDMDBlockRunner:
         """Request mode: all probes of one block, then its commit."""
         block_shape = (1, self.transformer.config.out_channels, *condition.shape[2:5])
         current_latents = randn_tensor(block_shape, generator=generator, device=self.device, dtype=torch.float32)
+        # Every forward below -- each probe and the commit -- reads the same camera trajectory, so the camera
+        # injector is built once per block: one cache for the whole loop, exactly as the stepwise path holds
+        # one per block across its separate calls.
+        camera_cache = CameraModulationCache()
         for step_index, (timestep_value, sigma) in enumerate(schedule):
             flow_prediction = self.probe_step(
                 current_latents=current_latents,
@@ -221,6 +235,7 @@ class LingBotDMDBlockRunner:
                 start_frame=start_frame,
                 timestep_value=timestep_value,
                 step_index=step_index,
+                camera_cache=camera_cache,
             )
             next_sigma = schedule[step_index + 1][1] if step_index + 1 < len(schedule) else None
             current_latents = self.apply_transition(
@@ -235,5 +250,6 @@ class LingBotDMDBlockRunner:
             cache=cache,
             ar=ar,
             start_frame=start_frame,
+            camera_cache=camera_cache,
         )
         return current_latents

vllm_omni/diffusion/models/lingbot_world/pipeline.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import functools
 import math
 import os
 from collections import OrderedDict
@@ -54,6 +55,7 @@ from vllm_omni.diffusion.models.lingbot_world.camera import (
 )
 from vllm_omni.diffusion.models.lingbot_world.dmd_block import ARBlockContext, LingBotDMDBlockRunner
 from vllm_omni.diffusion.models.lingbot_world.transformer import (
+    CameraModulationCache,
     CausalLingBotWorldTransformer3DModel,
     LingBotAttentionCache,
     LingBotTransformerCache,
@@ -91,7 +93,6 @@ logger = init_logger(__name__)
 _STREAMING_DECODE_BYTES_PER_PIXEL_FP32 = 37832 * 1024 / (64 * 64)
 
 if TYPE_CHECKING:
-    from diffusers.video_processor import VideoProcessor
     from tqdm.std import tqdm as TqdmProgressBar
 
     from vllm_omni.diffusion.worker.input_batch import InputBatch
@@ -152,6 +153,9 @@ class _LingBotRequestInputs:
     width: int
     num_frames: int
     num_latent_frames: int
+    # Request mode returns pixels through VideoProcessor in this type. Step execution (realtime sessions)
+    # returns every pixel type as ``(F, H, W, 3)`` uint8 frames, the layout the streaming endpoint muxes; only
+    # ``"latent"`` differs there. See ``_decode_chunk_to_pixels``.
     output_type: str
     max_sequence_length: int
     flow_shift: float
@@ -163,10 +167,20 @@ class _LingBotARSessionState:
     """Model-owned state; attention tensors remain runner-owned."""
 
     next_chunk_index: int = 0
-    prompt: str | None = None
+    # The umT5 embedding itself is reused by ``encode_prompt``'s own cache.
+    # ``text_cache_key`` does the same job for the runner-owned cross-attention
+    # K/V built from that embedding: it holds the (prompt, max_sequence_length,
+    # dtype) the K/V was built from and is only published after a successful
+    # build; ``text_cache_dirty`` is raised for the duration of a build, so a
+    # build that fails forces the next tick to invalidate whatever its prompt is.
+    text_cache_key: tuple[str, int, str] | None = None
+    text_cache_dirty: bool = False
     generator_state: torch.Tensor | None = None
     encoder_cache: list[torch.Tensor | None] | None = None
     pending_encoder_cache: list[torch.Tensor | None] | None = None
+    # The condition the encoder settles on once its history has stopped
+    # changing; see ``_next_condition_chunk``.
+    condition_fixed_point: torch.Tensor | None = None
     camera_tail: CameraTrajectory | None = None
     camera_pitch: float = 0.0
 
@@ -284,7 +298,70 @@ def _decode_source_image(image: PIL.Image.Image) -> PIL.Image.Image:
         raise ValueError(_SOURCE_IMAGE_ERROR) from None
 
 
+def _source_image_fingerprint(path: str | os.PathLike[str]) -> tuple[str, int, int, int, int, int]:
+    """Identify the bytes behind ``path``: the resolved path plus the stat fields that change when it is rewritten."""
+    resolved = os.path.realpath(os.fspath(path))
+    stat = os.stat(resolved)
+    return resolved, stat.st_dev, stat.st_ino, stat.st_size, stat.st_mtime_ns, stat.st_ctime_ns
+
+
+def _encoder_caches_identical(
+    before: list[torch.Tensor | None] | None, after: list[torch.Tensor | None] | None
+) -> bool:
+    """Whether two encoder histories hold the same tensors bit for bit, slot by slot."""
+    if before is None or after is None or len(before) != len(after):
+        return False
+    for x, y in zip(before, after, strict=True):
+        if x is None or y is None:
+            if x is not y:
+                return False
+        elif x.shape != y.shape or x.dtype != y.dtype or not torch.equal(x, y):
+            return False
+    return True
+
+
+def _uint8_frames(video: torch.Tensor) -> np.ndarray:
+    """Turn one decoded block ``(1, 3, F, H, W)`` in ``[-1, 1]`` into ``(F, H, W, 3)`` uint8 on the host.
+
+    This is the arithmetic ``VideoProcessor.postprocess_video`` applies for
+    ``"np"`` and ``"pil"`` -- denormalise, widen to float32, scale by 255 and
+    round -- done on the device before the copy instead of on the host after
+    it. The bytes are identical to the PIL path; what changes is that the
+    frames cross to the host once, already in the layout and dtype the
+    streaming endpoint muxes, instead of as a float32 copy that is then
+    converted per frame while the next block waits.
+    """
+    # Same order of operations as VideoProcessor: denormalise and clamp in
+    # the decoder dtype, then widen, scale and round.
+    frames = (video[0] / 2 + 0.5).clamp(0, 1).permute(1, 2, 3, 0)
+    frames = frames.float().mul_(255).round_().to(torch.uint8)
+    return np.ascontiguousarray(frames.cpu().numpy())
+
+
+@functools.lru_cache(maxsize=1)
+def _decode_source_image_fingerprinted(fingerprint: tuple[str, int, int, int, int, int]) -> PIL.Image.Image:
+    """The decoded image behind ``fingerprint`` (resolved path plus the stat fields that change on rewrite).
+
+    One entry: a request-mode caller that replays the same source path decodes it once. The realtime tick
+    carries its image inline and never comes through here.
+    """
+    return _decode_source_image_file(fingerprint[0])
+
+
 def _load_source_image(path: str | os.PathLike[str]) -> PIL.Image.Image:
+    """Decode the source image at ``path``, reusing the last decode while the file is unchanged.
+
+    Every caller gets its own copy, so a request that mutates its image cannot reach the next one. A path
+    that cannot be stat'ed goes straight to the uncached decode, which owns the validation error contract.
+    """
+    try:
+        fingerprint = _source_image_fingerprint(path)
+    except (OSError, TypeError, ValueError):
+        return _decode_source_image_file(path)
+    return _decode_source_image_fingerprinted(fingerprint).copy()
+
+
+def _decode_source_image_file(path: str | os.PathLike[str]) -> PIL.Image.Image:
     try:
         source_image = PIL.Image.open(path)
     except (OSError, SyntaxError, ValueError, PIL.Image.DecompressionBombError):
@@ -483,6 +560,9 @@ class LingBotWorldCausalDMDPipeline(
     """LingBot-World v2 I2V generation with a request-local causal cache."""
 
     supports_step_execution: ClassVar[bool] = True
+    # A realtime block is four probes and a commit over one prepared chunk
+    # state; nothing between those steps needs the serving scheduler.
+    supports_chunk_step_grouping: ClassVar[bool] = True
     _dit_modules: ClassVar[list[str]] = ["transformer"]
     _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
     _vae_modules: ClassVar[list[str]] = ["vae"]
@@ -637,12 +717,21 @@ class LingBotWorldCausalDMDPipeline(
             raise ValueError("LingBot AR-Diffusion cache needs a positive recent window after reserving sink frames.")
         condition_latent_frames = int(self.transformer.config.num_frames_per_block)
         condition_channels = self.vae_scale_factor_temporal + int(self.transformer.config.out_channels)
+        element_size = torch.empty((), dtype=self.transformer.dtype).element_size()
         condition_bytes_per_session = (
-            condition_channels
-            * condition_latent_frames
-            * latent_height
-            * latent_width
-            * torch.empty((), dtype=self.transformer.dtype).element_size()
+            condition_channels * condition_latent_frames * latent_height * latent_width * element_size
+        )
+        # Each layer retains scale and shift for the rank-local camera tokens
+        # across all forwards of a block, including between scheduler calls.
+        camera_tokens = tokens_per_frame * condition_latent_frames
+        sp_size = int(self.od_config.parallel_config.ulysses_degree)
+        camera_bytes_per_session = (
+            2
+            * int(self.transformer.config.num_layers)
+            * ((camera_tokens + sp_size - 1) // sp_size)
+            * int(self.transformer.config.num_attention_heads)
+            * int(self.transformer.config.attention_head_dim)
+            * element_size
         )
         return ARDiffusionKVCacheSpec(
             num_layers=int(self.transformer.config.num_layers),
@@ -658,10 +747,14 @@ class LingBotWorldCausalDMDPipeline(
                 ARDiffusionCrossAttentionKVSpec(
                     self._AR_TEXT_CACHE,
                     _MAX_SEQUENCE_LENGTH,
+                    # Cross-attention keeps every local head on every rank rather than sharding heads across
+                    # the Ulysses group, so the text pool holds that many heads, not the self-attention share.
+                    num_kv_heads=int(self.transformer.blocks[0].cross_attn.num_local_heads),
                 ),
             ),
             model_owned_state_bytes_per_session=(
                 condition_bytes_per_session
+                + camera_bytes_per_session
                 + 2 * self._condition_encoder_cache_bytes()  # Committed and in-flight histories.
                 + self._streaming_decode_bytes_per_session()
             ),
@@ -733,6 +826,7 @@ class LingBotWorldCausalDMDPipeline(
         if state is not None:
             state.encoder_cache = None
             state.pending_encoder_cache = None
+            state.condition_fixed_point = None
 
     def reset_ar_diffusion_session(self, session_id: str) -> None:
         self._release_condition_encoder_state(session_id)
@@ -988,6 +1082,50 @@ class LingBotWorldCausalDMDPipeline(
             )
         return condition
 
+    def _next_condition_chunk(
+        self,
+        inputs: _LingBotRequestInputs,
+        *,
+        start_frame: int,
+        session_state: _LingBotARSessionState,
+        dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, list[torch.Tensor | None]]:
+        """The session's condition for the block at ``start_frame``, encoding only while it still changes.
+
+        After a session's opening block the encoder is fed the same all-zero
+        pixel frames on every block, so the only thing that can make one
+        block's condition differ from the last is the encoder's own history.
+        The history is a causal cache with finite temporal reach, and once the
+        zero frames have flushed the source image out of it the encoder is at
+        a fixed point: identical input and identical state produce the same
+        output and the same state on every later block. That fixed point is
+        detected, not assumed after some number of blocks: a block whose
+        returned history is bit-identical to the history it was given is the
+        proof, and from then on the stored condition is returned and the
+        encoder is not run. Measured at 480x832 the encode is 178 ms of a
+        block and the fixed point arrives by the eleventh; the served
+        rollouts this was measured on (40 blocks) are then byte-identical.
+
+        The opening block resets the fixed point, so a session that restarts
+        from a new source image encodes again until it settles.
+        """
+        if start_frame == 0:
+            session_state.condition_fixed_point = None
+        elif session_state.condition_fixed_point is not None:
+            assert session_state.encoder_cache is not None
+            # The history is unchanged by construction; the pending copy is
+            # what post_decode commits, exactly as an encoded block's would be.
+            return session_state.condition_fixed_point.clone(), list(session_state.encoder_cache)
+        condition, pending_cache = self._prepare_condition_chunk(
+            inputs,
+            start_frame=start_frame,
+            encoder_cache=session_state.encoder_cache,
+            dtype=dtype,
+        )
+        if start_frame > 0 and _encoder_caches_identical(session_state.encoder_cache, pending_cache):
+            session_state.condition_fixed_point = condition.clone()
+        return condition, pending_cache
+
     @torch.no_grad()
     def _prepare_condition_chunk(
         self,
@@ -1160,10 +1298,8 @@ class LingBotWorldCausalDMDPipeline(
                 for block in self.transformer.blocks:
                     cross_attention = block.cross_attn
                     shape = (cross_attention.num_local_heads, cross_attention.head_dim)
-                    key = cross_attention.shard_kv_heads(
-                        cross_attention.norm_k(cross_attention.k(projected_text)).unflatten(2, shape)
-                    )
-                    value = cross_attention.shard_kv_heads(cross_attention.v(projected_text).unflatten(2, shape))
+                    key = cross_attention.norm_k(cross_attention.k(projected_text)).unflatten(2, shape)
+                    value = cross_attention.v(projected_text).unflatten(2, shape)
                     yield key, value
 
             state.populate_cross_attention(
@@ -1243,6 +1379,34 @@ class LingBotWorldCausalDMDPipeline(
             progress_bar=progress_bar,
         )
 
+    def _session_text_caches(
+        self,
+        prompt_embeds: torch.Tensor,
+        *,
+        prompt: str,
+        max_sequence_length: int,
+        session_state: _LingBotARSessionState,
+    ) -> list[LingBotAttentionCache]:
+        """Runner-owned cross-attention K/V for this tick, rebuilt when its inputs change.
+
+        The K/V is keyed by everything that defines the prompt embedding
+        (text, sequence length, dtype). The first build of a session needs no
+        invalidation (the runner state is empty); a key change does. The dirty
+        flag is raised for the duration of the build and cleared only after
+        ``_ar_text_caches`` returns, so a build that fails after invalidating
+        or partially writing the caches makes the next tick, whatever its
+        prompt, invalidate again instead of consuming partial K/V.
+        """
+        text_key = (prompt, int(max_sequence_length), str(prompt_embeds.dtype))
+        previous = session_state.text_cache_key
+        invalidate = session_state.text_cache_dirty or (previous is not None and previous != text_key)
+        session_state.text_cache_dirty = True
+        session_state.text_cache_key = None
+        caches = self._ar_text_caches(prompt_embeds, invalidate=invalidate)
+        session_state.text_cache_key = text_key
+        session_state.text_cache_dirty = False
+        return caches
+
     def encode_prompt(
         self,
         prompt: str,
@@ -1345,6 +1509,9 @@ class LingBotWorldCausalDMDPipeline(
         )
         dtype = self.transformer.dtype
         # Phase 1: turn all three user inputs into DiT-ready conditions.
+        # ``encode_prompt`` reuses an already-encoded prompt itself; the
+        # runner-owned cross-attention K/V is keyed the same way in
+        # ``_session_text_caches`` below.
         prompt_embeds = self.encode_prompt(
             inputs.prompt,
             max_sequence_length=inputs.max_sequence_length,
@@ -1354,10 +1521,10 @@ class LingBotWorldCausalDMDPipeline(
             condition = self._prepare_condition(inputs, dtype=dtype)
         else:
             assert session_state is not None
-            condition, pending_encoder_cache = self._prepare_condition_chunk(
+            condition, pending_encoder_cache = self._next_condition_chunk(
                 inputs,
                 start_frame=tick.chunk_index * block_frames,
-                encoder_cache=session_state.encoder_cache,
+                session_state=session_state,
                 dtype=dtype,
             )
         camera_pitch: float | None = None
@@ -1396,12 +1563,13 @@ class LingBotWorldCausalDMDPipeline(
             )
         else:
             assert session_state is not None
-            prompt_changed = session_state.prompt is not None and session_state.prompt != inputs.prompt
             if session_state.generator_state is not None:
                 inputs.generator.set_state(session_state.generator_state)
-            ar_cross_attention = self._ar_text_caches(
+            ar_cross_attention = self._session_text_caches(
                 prompt_embeds,
-                invalidate=prompt_changed,
+                prompt=inputs.prompt,
+                max_sequence_length=inputs.max_sequence_length,
+                session_state=session_state,
             )
         generated_blocks: list[torch.Tensor] = []
         total_steps = (inputs.num_latent_frames // block_frames) * len(LINGBOT_DMD_TIMESTEPS)
@@ -1437,7 +1605,6 @@ class LingBotWorldCausalDMDPipeline(
         cache = None
         if tick is not None:
             assert session_state is not None
-            session_state.prompt = inputs.prompt
             session_state.generator_state = inputs.generator.get_state()
             session_state.encoder_cache = pending_encoder_cache
             session_state.camera_tail = camera_tail
@@ -1536,11 +1703,15 @@ class LingBotWorldCausalDMDPipeline(
         if state is not None:
             state.release()
 
-    def _decode_chunk_to_pixels(
-        self, latents: torch.Tensor, *, output_type: str, session_id: str | None = None
-    ) -> torch.Tensor | np.ndarray | list[list[PIL.Image.Image]]:
+    def _decode_chunk_to_pixels(self, latents: torch.Tensor, *, session_id: str | None = None) -> np.ndarray:
         """Decode one AR block so streaming consumers receive pixels, not latents.
 
+        The block comes back as ``(F, H, W, 3)`` uint8 whatever pixel output
+        type the request named: the realtime endpoint muxes uint8 frames, so
+        ``"pil"`` and ``"np"`` only ever differed in how much work stood
+        between the decoder and those bytes. Only ``"latent"`` means
+        something else, and the caller keeps latents before reaching here.
+
         With a ``session_id`` the block is decoded through that session's own
         temporal cache, so chunk *N + 1* continues chunk *N*: the causal decoder
         expands the session's opening latent frame to one raw frame and every
@@ -1566,7 +1737,7 @@ class LingBotWorldCausalDMDPipeline(
             video = self.vae.decode(vae_latents, return_dict=False)[0]
         else:
             video = self._streaming_decode_chunk(decoder, vae_latents, cast(str, session_id))
-        return self._video_processor().postprocess_video(video, output_type=output_type)
+        return _uint8_frames(video)
 
     def _streaming_decode_chunk(
         self, decoder: WanStreamingDecoder, vae_latents: torch.Tensor, session_id: str
@@ -1587,15 +1758,6 @@ class LingBotWorldCausalDMDPipeline(
             self._release_streaming_decode_state(session_id)
             raise
 
-    def _video_processor(self) -> VideoProcessor:
-        processor = getattr(self, "_cached_video_processor", None)
-        if processor is None:
-            from diffusers.video_processor import VideoProcessor
-
-            processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
-            self._cached_video_processor = processor
-        return processor
-
     def _require_bound_ar_state(self) -> None:
         if self._ar_diffusion_kv_state is None:
             raise RuntimeError("LingBot step execution requires AR-Diffusion session binding.")
@@ -1683,10 +1845,10 @@ class LingBotWorldCausalDMDPipeline(
         session_state = self._ar_sessions[state.request_id]
         if session_state.next_chunk_index != state.chunk_index:
             raise ValueError("LingBot condition chunks must be contiguous within the session.")
-        condition, pending_encoder_cache = self._prepare_condition_chunk(
+        condition, pending_encoder_cache = self._next_condition_chunk(
             inputs,
             start_frame=start_frame,
-            encoder_cache=session_state.encoder_cache,
+            session_state=session_state,
             dtype=extra["dtype"],
         )
         previous = extra.get("camera_tail")
@@ -1780,6 +1942,12 @@ class LingBotWorldCausalDMDPipeline(
             raise ValueError("LingBot World requires the runner-provided torch.Generator.")
         extra["condition"] = condition
         extra["camera"] = camera
+        # One camera-modulation cache per block: the four probes and the
+        # commit below are separate scheduler steps, so the block runner
+        # cannot open the reuse window around them the way generate_block
+        # does around its own loop. Holding the cache here makes the window
+        # span exactly this block, per request.
+        extra["camera_cache"] = CameraModulationCache()
         extra["camera_tail"] = camera_tail
         extra["start_frame"] = start_frame
         extra["ar_cross_attention"] = self._ar_text_caches(
@@ -1838,6 +2006,7 @@ class LingBotWorldCausalDMDPipeline(
             start_frame=int(extra["start_frame"]),
             timestep_value=float(schedule[step_in_chunk][0]),
             step_index=step_in_chunk,
+            camera_cache=extra["camera_cache"],
         )
 
     def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None:
@@ -1880,6 +2049,7 @@ class LingBotWorldCausalDMDPipeline(
             cache=None,
             ar=self._ar_block_context(extra["ar_cross_attention"]),
             start_frame=int(extra["start_frame"]),
+            camera_cache=extra.pop("camera_cache"),
         )
         completed_chunk_index = state.chunk_index
         inputs: _LingBotRequestInputs = extra["inputs"]
@@ -1889,13 +2059,7 @@ class LingBotWorldCausalDMDPipeline(
         if inputs.output_type == "latent":
             payload: dict[str, Any] = {"latents": latents}
         else:
-            payload = {
-                "video": self._decode_chunk_to_pixels(
-                    latents,
-                    output_type=inputs.output_type,
-                    session_id=state.request_id,
-                )
-            }
+            payload = {"video": self._decode_chunk_to_pixels(latents, session_id=state.request_id)}
         output = {
             "payload": payload,
             "metadata": {

vllm_omni/diffusion/models/lingbot_world/transformer.py

@@ -24,7 +24,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.utils import set_weight_attrs
 
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D
+from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D, all_to_all_5D
 from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
 from vllm_omni.diffusion.layers.norm import LayerNorm
@@ -329,9 +329,18 @@ class LingBotSelfAttention(nn.Module):
             key = self.rotary_embedding(key, cos, sin)
 
         if self.ulysses_world_size > 1:
-            query = SeqAllToAll4D.apply(self.ulysses_group, query, 2, 1, False)
-            key = SeqAllToAll4D.apply(self.ulysses_group, key, 2, 1, False)
-            value = SeqAllToAll4D.apply(self.ulysses_group, value, 2, 1, False)
+            # One fused sequence->head all-to-all for q/k/v instead of three:
+            # (B, S/N, 3, H, D) -> (B, S, 3, H/N, D). Same head-group assignment
+            # as the per-tensor 4D exchange (rank r receives head group r), so
+            # the result is a pure permutation of the three-call version, with a
+            # 3x larger message and one third of the latency-bound collectives.
+            qkv = all_to_all_5D(
+                torch.stack((query, key, value), dim=2),
+                scatter_idx=3,
+                gather_idx=1,
+                group=self.ulysses_group,
+            )
+            query, key, value = qkv.unbind(2)
 
         if isinstance(cache, ARDiffusionPagedLayerInputs):
             if query.shape[0] != 1:
@@ -421,13 +430,11 @@ class LingBotCrossAttention(nn.Module):
         self.num_heads = num_heads
         self.head_dim = dim // num_heads
         self.num_local_heads = num_heads // tp_size
-        self.ulysses_world_size, self.ulysses_rank, self.ulysses_group = _ulysses_state()
-        if self.num_local_heads % self.ulysses_world_size:
-            raise ValueError(
-                "LingBot local cross-attention heads must be divisible by the Ulysses degree: "
-                f"heads={self.num_local_heads}, ulysses={self.ulysses_world_size}."
-            )
-        self.num_sp_heads = self.num_local_heads // self.ulysses_world_size
+        # Cross-attention keeps every local head on every Ulysses rank and attends the rank's own sequence
+        # shard. Each query token attends the whole text sequence independently of the others, so there is
+        # nothing to exchange: the sequence split Ulysses already gives us is the only split needed, and both
+        # all-to-all collectives -- one for the query, one for the output -- disappear. The cost is that the text K/V is
+        # replicated per rank instead of head-sharded; ar_diffusion_kv_cache_spec sizes the pool for it.
         self.tp_inner_dim = self.num_local_heads * self.head_dim
 
         self.q = ColumnParallelLinear(
@@ -465,9 +472,9 @@ class LingBotCrossAttention(nn.Module):
         self.norm_q = _LingBotRMSNorm(self.tp_inner_dim, eps)
         self.norm_k = _LingBotRMSNorm(self.tp_inner_dim, eps)
         self.attn = Attention(
-            num_heads=self.num_sp_heads,
+            num_heads=self.num_local_heads,
             head_size=self.head_dim,
-            num_kv_heads=self.num_sp_heads,
+            num_kv_heads=self.num_local_heads,
             softmax_scale=self.head_dim**-0.5,
             causal=False,
             role="cross",
@@ -477,14 +484,6 @@ class LingBotCrossAttention(nn.Module):
             disable_kv_quant=True,
         )
 
-    def shard_kv_heads(self, value: torch.Tensor) -> torch.Tensor:
-        """Select this Ulysses rank's projected K/V heads."""
-
-        if self.ulysses_world_size == 1:
-            return value
-        start = self.ulysses_rank * self.num_sp_heads
-        return value[:, :, start : start + self.num_sp_heads].clone(memory_format=torch.contiguous_format)
-
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -494,8 +493,6 @@ class LingBotCrossAttention(nn.Module):
     ) -> tuple[torch.Tensor, LingBotAttentionCache]:
         query = self.norm_q(self.q(hidden_states))
         query = query.unflatten(2, (self.num_local_heads, self.head_dim))
-        if self.ulysses_world_size > 1:
-            query = SeqAllToAll4D.apply(self.ulysses_group, query, 2, 1, False)
 
         # Text K/V is constant within a request and is projected once per layer.
         if cache is None:
@@ -503,8 +500,8 @@ class LingBotCrossAttention(nn.Module):
                 raise ValueError("encoder_hidden_states are required when the cross-attention cache is empty.")
             key = self.norm_k(self.k(encoder_hidden_states))
             value = self.v(encoder_hidden_states)
-            key = self.shard_kv_heads(key.unflatten(2, (self.num_local_heads, self.head_dim)))
-            value = self.shard_kv_heads(value.unflatten(2, (self.num_local_heads, self.head_dim)))
+            key = key.unflatten(2, (self.num_local_heads, self.head_dim))
+            value = value.unflatten(2, (self.num_local_heads, self.head_dim))
             cache = LingBotAttentionCache(
                 key=key,
                 value=value,
@@ -517,8 +514,6 @@ class LingBotCrossAttention(nn.Module):
             value = cache.value[:, : cache.end]
 
         output = self.attn(query, key, value)
-        if self.ulysses_world_size > 1:
-            output = SeqAllToAll4D.apply(self.ulysses_group, output, 1, 2, False)
         return self.o(output.flatten(2, 3)), cache
 
 
@@ -592,6 +587,16 @@ class LingBotAttentionBlock(nn.Module):
         self.cam_scale_layer = nn.Linear(dim, dim)
         self.cam_shift_layer = nn.Linear(dim, dim)
 
+    def camera_modulation(self, camera_hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Build this block's camera scale and shift from the patchified camera tokens.
+
+        Depends only on ``camera_hidden_states``, so a caller holding the camera trajectory fixed -- every
+        denoise forward of one AR chunk does -- can build this once and pass it back through ``forward``.
+        """
+        camera_features = self.cam_injector_layer2(F.silu(self.cam_injector_layer1(camera_hidden_states)))
+        camera_features = camera_features + camera_hidden_states
+        return self.cam_scale_layer(camera_features), self.cam_shift_layer(camera_features)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -605,6 +610,7 @@ class LingBotAttentionBlock(nn.Module):
         sink_tokens: int,
         update_cache: bool,
         rotary_emb: tuple[torch.Tensor, torch.Tensor],
+        camera_modulation: tuple[torch.Tensor, torch.Tensor] | None = None,
     ) -> tuple[torch.Tensor, LingBotAttentionCache]:
         batch_size, token_count, dim = hidden_states.shape
         if (
@@ -638,10 +644,9 @@ class LingBotAttentionBlock(nn.Module):
         hidden_grid = hidden_grid + attention_output.unflatten(1, (num_frames, tokens_per_frame)) * gate_msa
         hidden_states = hidden_grid.flatten(1, 2).to(hidden_states.dtype)
 
-        camera_features = self.cam_injector_layer2(F.silu(self.cam_injector_layer1(camera_hidden_states)))
-        camera_features = camera_features + camera_hidden_states
-        camera_scale = self.cam_scale_layer(camera_features)
-        camera_shift = self.cam_shift_layer(camera_features)
+        camera_scale, camera_shift = (
+            self.camera_modulation(camera_hidden_states) if camera_modulation is None else camera_modulation
+        )
         hidden_states = ((1 + camera_scale) * hidden_states + camera_shift).to(hidden_states.dtype)
 
         attention_output, cross_cache = self.cross_attn(
@@ -763,6 +768,22 @@ def _rope_axis(max_seq_len: int, dim: int, *, start: int = 0) -> tuple[torch.Ten
     return phase.cos().float(), phase.sin().float()
 
 
+class CameraModulationCache:
+    """Per-block camera scale and shift, held by the caller across the forwards of one AR block.
+
+    The camera injector is four projections and a SiLU per block and reads only the camera tokens, which do
+    not change between the denoise steps and the clean commit of one block. The caller that owns that
+    invariant creates one of these per block and passes it to every forward of the block; ``forward`` fills a
+    block's entry on first use and reads it after. Its lifetime is the caller's: the AR runner keeps one per
+    block, ``generate_block`` one per loop, and nothing here checks the camera tensor for change.
+    """
+
+    __slots__ = ("entries",)
+
+    def __init__(self) -> None:
+        self.entries: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
+
+
 class CausalLingBotWorldTransformer3DModel(nn.Module):
     """Checkpoint-compatible causal LingBot World video transformer."""
 
@@ -922,6 +943,27 @@ class CausalLingBotWorldTransformer3DModel(nn.Module):
             cosine, sine = _rope_axis(rope_max_seq_len, axis_dim)
             self.register_buffer(f"_rope_{axis}_cosine", cosine, persistent=False)
             self.register_buffer(f"_rope_{axis}_sine", sine, persistent=False)
+        # One static-address buffer per shape for the per-token timestep
+        # projection; see _stage_timestep_projection.
+        self._timestep_projection_buffers: dict[tuple[tuple[int, ...], torch.dtype, torch.device], torch.Tensor] = {}
+
+    def _stage_timestep_projection(self, projection: torch.Tensor) -> torch.Tensor:
+        """Hand the blocks the timestep projection from a buffer whose address never changes.
+
+        Every regional block graph takes the projection as an input. Under CUDA-graph replay an input that
+        arrives in a fresh tensor each forward is copied into the graph's placeholder on every replay: at
+        480x832 with four Ulysses ranks that is a 72 MB copy per block, forty times per forward, two hundred
+        times per AR chunk. Copying the projection once per forward into a buffer marked static lets every
+        replay read it in place. The values the blocks see are the same; only where they live changes.
+        """
+        key = (tuple(projection.shape), projection.dtype, projection.device)
+        buffer = self._timestep_projection_buffers.get(key)
+        if buffer is None:
+            buffer = torch.empty(projection.shape, dtype=projection.dtype, device=projection.device)
+            torch._dynamo.mark_static_address(buffer)
+            self._timestep_projection_buffers[key] = buffer
+        buffer.copy_(projection)
+        return buffer
 
     @property
     def dtype(self) -> torch.dtype:
@@ -1178,6 +1220,7 @@ class CausalLingBotWorldTransformer3DModel(nn.Module):
         cache: LingBotTransformerCache,
         start_frame: int,
         update_cache: bool,
+        camera_modulation_cache: CameraModulationCache | None = None,
     ) -> torch.Tensor:
         if not isinstance(update_cache, bool):
             raise ValueError("update_cache must be a boolean.")
@@ -1272,16 +1315,25 @@ class CausalLingBotWorldTransformer3DModel(nn.Module):
                     "LingBot SP input hooks did not shard all conditioning tensors consistently; "
                     f"expected {local_tokens} tokens per rank. Check SP hook registration and input dimensions."
                 )
+        timestep_projection = self._stage_timestep_projection(timestep_projection)
         rotary_emb = (cosine, sine)
         # Phase 3: each layer receives its own cache entry. Text K/V is passed
         # only when absent; the returned cache is stored for subsequent DMD
-        # steps and causal blocks in this request.
+        # steps and causal blocks in this request. A caller-held camera cache
+        # gives every block its modulation once per AR block.
         for index, block in enumerate(self.blocks):
+            camera_modulation = None
+            if camera_modulation_cache is not None:
+                camera_modulation = camera_modulation_cache.entries.get(index)
+                if camera_modulation is None:
+                    camera_modulation = block.camera_modulation(camera_hidden_states)
+                    camera_modulation_cache.entries[index] = camera_modulation
             hidden_states, cross_cache = block(
                 hidden_states,
                 projected_text if cache.cross_attention[index] is None else None,
                 timestep_projection,
                 camera_hidden_states,
+                camera_modulation=camera_modulation,
                 self_cache=cache.self_attention[index],
                 cross_cache=cache.cross_attention[index],
                 current_start=current_start,