From 5e6f49c986b5e629a272caeb6963b088fdcf69b0 Mon Sep 17 00:00:00 2001 From: HuangJi <32611516+IPostYellow@users.noreply.github.com> Date: Sun, 5 Jul 2026 11:11:54 +0800 Subject: [PATCH] [diffusion] feat: add LingBot realtime prompt, KV window, and lazy VAE controls (#30040) --- .../configs/pipeline_configs/lingbot_world.py | 67 ++- python/sglang/multimodal_gen/envs.py | 8 + .../lingbot_world_realtime_adapter.py | 173 ++++++- .../layers/kvcache/causal_attention_cache.py | 121 ++++- .../runtime/models/dits/lingbot_world.py | 26 +- .../lingbot_world_causal_dmd_pipeline.py | 2 +- .../runtime/pipelines_core/schedule_batch.py | 1 + .../pipelines_core/stages/causal_denoising.py | 8 +- .../lingbot_world/__init__.py | 13 +- .../lingbot_world/constants.py | 11 + .../lingbot_world_causal_denoising.py | 383 +++++++++++++++- .../realtime/test_lingbot_causal_denoising.py | 428 +++++++++++++++++- .../unit/realtime/test_realtime_runtime.py | 102 ++++- 13 files changed, 1286 insertions(+), 57 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/constants.py diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/lingbot_world.py b/python/sglang/multimodal_gen/configs/pipeline_configs/lingbot_world.py index 49b81b416..fbaf6644e 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/lingbot_world.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/lingbot_world.py @@ -11,6 +11,7 @@ from typing import Any import numpy as np import torch +from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.models import DiTConfig from sglang.multimodal_gen.configs.models.dits import LingBotWorldVideoConfig from sglang.multimodal_gen.configs.pipeline_configs.wan import Wan2_2_I2V_A14B_Config @@ -325,6 +326,44 @@ class LingBotWorldCausalDMDConfig(LingBotWorldI2VConfig): warp_denoising_step: bool = True realtime_causal_sink_size: int | None = None realtime_causal_kv_cache_num_frames: int | None = None + interactive_kv_window_enable: bool = False + interactive_kv_still_window: int | None = 3 + interactive_kv_moving_window: int | None = 12 + interactive_kv_still_chunks: int = 2 + lazy_vae_encode_black_frames: int = 0 + + def preprocess_vae_encode(self, image, vae): + image = super().preprocess_vae_encode(image, vae) + lazy_black_frames = envs.SGLANG_LINGBOT_LAZY_VAE_ENCODE_BLACK_FRAMES + if lazy_black_frames is None: + lazy_black_frames = self.lazy_vae_encode_black_frames + lazy_black_frames = max(0, int(lazy_black_frames or 0)) + if lazy_black_frames <= 0 or image.ndim != 5: + return image + + num_frames = int(image.shape[2]) + if num_frames <= 1: + return image + + temporal_ratio = int(self.vae_config.arch_config.temporal_compression_ratio) + encode_frames = min(num_frames, 1 + lazy_black_frames) + if (encode_frames - 1) % temporal_ratio != 0: + encode_frames = (encode_frames - 1) // temporal_ratio + 1 + encode_frames = encode_frames * temporal_ratio + 1 + encode_frames = min(num_frames, encode_frames) + + if encode_frames >= num_frames: + return image + + logger.info( + "LingBot lazy VAE encode: pixel_frames=%s encode_pixel_frames=%s " + "black_frames=%s temporal_ratio=%s", + num_frames, + encode_frames, + lazy_black_frames, + temporal_ratio, + ) + return image[:, :, :encode_frames].contiguous() def postprocess_image_latent(self, latent_condition, batch): """Build condition tensor aligned to chunk_size (num_frames_per_block). @@ -342,9 +381,31 @@ class LingBotWorldCausalDMDConfig(LingBotWorldI2VConfig): latent_width = batch.width // spatial_ratio # Align num_latent_frames to chunk_size - num_latent_frames = latent_condition.shape[2] - num_latent_frames = num_latent_frames - (num_latent_frames % chunk_size) - latent_condition = latent_condition[:, :, :num_latent_frames, :, :] + target_latent_frames = (int(batch.num_frames) - 1) // temporal_ratio + 1 + target_latent_frames = target_latent_frames - ( + target_latent_frames % chunk_size + ) + encoded_latent_frames = int(latent_condition.shape[2]) + if encoded_latent_frames < target_latent_frames: + tail = latent_condition[:, :, -1:, :, :].repeat( + 1, + 1, + target_latent_frames - encoded_latent_frames, + 1, + 1, + ) + latent_condition = torch.cat([latent_condition, tail], dim=2) + elif encoded_latent_frames > target_latent_frames: + latent_condition = latent_condition[:, :, :target_latent_frames, :, :] + + num_latent_frames = int(latent_condition.shape[2]) + if encoded_latent_frames != num_latent_frames: + logger.info( + "LingBot lazy VAE condition: encoded_latent_frames=%s " + "target_latent_frames=%s", + encoded_latent_frames, + num_latent_frames, + ) # Number of initial frames that have actual image content # (latent_condition from VAE encode of [image, zeros...]) diff --git a/python/sglang/multimodal_gen/envs.py b/python/sglang/multimodal_gen/envs.py index 5e5182521..72c89b2ec 100644 --- a/python/sglang/multimodal_gen/envs.py +++ b/python/sglang/multimodal_gen/envs.py @@ -57,6 +57,8 @@ if TYPE_CHECKING: SGLANG_CACHE_DIT_SECONDARY_TS_ORDER: int = 1 # model loading SGLANG_USE_RUNAI_MODEL_STREAMER: bool = True + SGLANG_LINGBOT_ENABLE_INTERACTIVE_KV_WINDOW: bool = False + SGLANG_LINGBOT_LAZY_VAE_ENCODE_BLACK_FRAMES: int | None = None SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND: str | None = None SGLANG_DIFFUSION_ENABLE_W8A8_FP8_GEMM: bool = False SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D: str = "auto" @@ -295,6 +297,12 @@ environment_variables: dict[str, Callable[[], Any]] = { "SGLANG_USE_RUNAI_MODEL_STREAMER": _lazy_bool( "SGLANG_USE_RUNAI_MODEL_STREAMER", "true" ), + "SGLANG_LINGBOT_ENABLE_INTERACTIVE_KV_WINDOW": _lazy_bool( + "SGLANG_LINGBOT_ENABLE_INTERACTIVE_KV_WINDOW" + ), + "SGLANG_LINGBOT_LAZY_VAE_ENCODE_BLACK_FRAMES": _lazy_int( + "SGLANG_LINGBOT_LAZY_VAE_ENCODE_BLACK_FRAMES" + ), # FlashInfer FP4 GEMM backend override for diffusion NVFP4. # When unset, diffusion ModelOpt NVFP4 defaults to flashinfer_trtllm. # Supported values: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/lingbot_world_realtime_adapter.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/lingbot_world_realtime_adapter.py index 048bd2349..5313459aa 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/lingbot_world_realtime_adapter.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/lingbot_world_realtime_adapter.py @@ -14,7 +14,15 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_adapter build_realtime_sampling_params, save_realtime_first_frame, ) -from sglang.multimodal_gen.runtime.realtime.control_signals import ControlSignalQueue +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world.constants import ( + LINGBOT_CAMERA_ACTIONS_CONDITION, + LINGBOT_PROMPT_UPDATED_CONDITION, +) +from sglang.multimodal_gen.runtime.realtime.control_signals import ( + ControlSignalQueue, + ParsedControlEventPayload, + parse_control_event_payload, +) from sglang.multimodal_gen.runtime.realtime.states import ( RealtimeCameraControlState, ) @@ -29,6 +37,7 @@ if TYPE_CHECKING: LINGBOT_REALTIME_DEFAULT_NUM_INFERENCE_STEPS = 4 LINGBOT_REALTIME_MIN_CONDITION_CHUNKS = 2 +COMPOSITE_INPUT_EVENT_KIND = "composite_input" class LingBotWorldRealtimeState(RealtimeCameraControlState): @@ -47,16 +56,44 @@ class LingBotWorldRealtimeState(RealtimeCameraControlState): def receive_prompt(self, prompt: str, *, event_id: int | None = None) -> None: self.prompt_queue.push("prompt", prompt, event_id=event_id) + def parse_camera_control_event_payload( + self, + payload: Any, + *, + event_id: int | None, + ) -> ParsedControlEventPayload: + return parse_control_event_payload( + payload, + event_id=event_id, + kind="camera_actions", + normalize_state_payload=self._normalize_state_actions, + validate_script_payload=LingBotWorldRealtimeAdapter._validate_camera_actions, + ) + + def receive_parsed_camera_control_event_payload( + self, + parsed: ParsedControlEventPayload, + *, + event_id: int | None, + ) -> str: + if parsed.mode == "state": + transitions = parsed.payload + self.receive_camera_state_transitions(transitions) + return f"kind=camera_actions, mode=state, transitions={len(transitions)}" + + camera_actions = parsed.payload + self.receive_camera_action_script(camera_actions, event_id=event_id) + return f"kind=camera_actions, mode=script, frames={len(camera_actions)}" + def receive_camera_control_event_payload( self, payload: Any, *, event_id: int | None, ) -> str: - return super().receive_camera_control_event_payload( - payload, - event_id=event_id, - validate_camera_actions=LingBotWorldRealtimeAdapter._validate_camera_actions, + parsed = self.parse_camera_control_event_payload(payload, event_id=event_id) + return self.receive_parsed_camera_control_event_payload( + parsed, event_id=event_id ) def sample_prompt(self) -> str: @@ -86,7 +123,7 @@ class LingBotWorldRealtimeAdapter(BaseRealtimeModelAdapter): request: RealtimeVideoGenerationsRequest, ) -> None: condition_inputs = request.condition_inputs or {} - camera_actions = condition_inputs.get("camera_actions") + camera_actions = condition_inputs.get(LINGBOT_CAMERA_ACTIONS_CONDITION) if camera_actions is not None: state = self._state(session) state.receive_camera_action_script( @@ -113,17 +150,119 @@ class LingBotWorldRealtimeAdapter(BaseRealtimeModelAdapter): ) -> str: state = self._state(session) if event.kind == "camera_actions": - return state.receive_camera_control_event_payload( - event.payload, - event_id=event.event_id, - ) + return self._ingest_camera_actions(state, event.payload, event.event_id) elif event.kind == "prompt": - if not isinstance(event.payload, str) or not event.payload: - raise ValueError("prompt event payload must be a non-empty string") - state.receive_prompt(event.payload, event_id=event.event_id) - return f"kind=prompt, prompt_len={len(event.payload)}" + return self._ingest_prompt(state, event.payload, event.event_id) + elif event.kind == COMPOSITE_INPUT_EVENT_KIND: + return self._ingest_composite_input(state, event.payload, event.event_id) raise ValueError(f"unsupported event kind: {event.kind}") + def _ingest_camera_actions( + self, + state: LingBotWorldRealtimeState, + payload: Any, + event_id: int | None, + ) -> str: + return state.receive_camera_control_event_payload( + payload, + event_id=event_id, + ) + + def _ingest_prompt( + self, + state: LingBotWorldRealtimeState, + payload: Any, + event_id: int | None, + ) -> str: + prompt = self._validate_prompt_payload(payload) + state.receive_prompt(prompt, event_id=event_id) + return f"kind=prompt, prompt_len={len(prompt)}" + + @staticmethod + def _validate_prompt_payload(payload: Any) -> str: + if not isinstance(payload, str) or not payload: + raise ValueError("prompt event payload must be a non-empty string") + return payload + + def _ingest_composite_input( + self, + state: LingBotWorldRealtimeState, + payload: Any, + event_id: int | None, + ) -> str: + if not isinstance(payload, dict): + raise ValueError("composite_input event payload must be a map") + input_types = payload.get("input_types") + if not isinstance(input_types, list) or not input_types: + raise ValueError( + "composite_input event payload requires non-empty input_types" + ) + + parsed_inputs = [] + for input_type in input_types: + if not isinstance(input_type, str) or not input_type: + raise ValueError( + "composite_input input_types must contain non-empty strings" + ) + if input_type not in payload: + raise ValueError(f"composite_input event payload requires {input_type}") + parsed_inputs.append( + ( + input_type, + self._parse_composite_input_item( + state, + input_type, + payload[input_type], + event_id, + ), + ) + ) + + input_logs = [] + for input_type, parsed_payload in parsed_inputs: + input_logs.append( + self._ingest_parsed_composite_input_item( + state, + input_type, + parsed_payload, + event_id, + ) + ) + return f"kind=composite_input, inputs={input_logs}" + + def _parse_composite_input_item( + self, + state: LingBotWorldRealtimeState, + input_type: str, + payload: Any, + event_id: int | None, + ) -> Any: + if input_type == "camera_actions": + return state.parse_camera_control_event_payload( + payload, + event_id=event_id, + ) + if input_type == "prompt": + return self._validate_prompt_payload(payload) + raise ValueError(f"unsupported composite_input type: {input_type}") + + def _ingest_parsed_composite_input_item( + self, + state: LingBotWorldRealtimeState, + input_type: str, + parsed_payload: Any, + event_id: int | None, + ) -> str: + if input_type == "camera_actions": + return state.receive_parsed_camera_control_event_payload( + parsed_payload, + event_id=event_id, + ) + if input_type == "prompt": + state.receive_prompt(parsed_payload, event_id=event_id) + return f"kind=prompt, prompt_len={len(parsed_payload)}" + raise ValueError(f"unsupported composite_input type: {input_type}") + def sample_chunk_inputs( self, session: GenerateSession, @@ -137,18 +276,22 @@ class LingBotWorldRealtimeAdapter(BaseRealtimeModelAdapter): if request is None: raise ValueError("realtime request is not initialized") + prompt_updated = False if chunk.index == 0: prompt = request.prompt elif state.has_prompt(): prompt = state.sample_prompt() request.prompt = prompt + prompt_updated = True else: prompt = request.prompt condition_inputs = {} + if prompt_updated: + condition_inputs[LINGBOT_PROMPT_UPDATED_CONDITION] = True camera_actions = state.sample_camera_actions(chunk_size) if camera_actions is not None: - condition_inputs["camera_actions"] = camera_actions + condition_inputs[LINGBOT_CAMERA_ACTIONS_CONDITION] = camera_actions return RealtimeChunkInputs(prompt=prompt, condition_inputs=condition_inputs) def build_sampling_params( diff --git a/python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py b/python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py index c61037338..2d6ae23d7 100644 --- a/python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py +++ b/python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py @@ -110,6 +110,7 @@ class CausalSelfAttentionKVCache: value: torch.Tensor, current_chunk_start: int, cache_head_start: int | None = None, + recent_window_tokens: int | None = None, debug_name: str = "causal KV cache", ) -> CausalAttentionKVView: """write fresh kv into the cache, returns the part of view visible to the current chunk @@ -118,6 +119,11 @@ class CausalSelfAttentionKVCache: current_chunk_start: the global position of the start of the chunk cache_head_start: first cache head for key/value when they only carry a local slice of the cache heads; other heads are left untouched + recent_window_tokens: recent-window attention size. ``None`` + returns the full visible attention window. ``0`` keeps only sink + tokens plus the current chunk. A positive value keeps sink tokens, + up to that many tokens before the current chunk, and the current + chunk. Negative values are invalid. """ num_new_tokens = key.shape[1] @@ -265,17 +271,22 @@ class CausalSelfAttentionKVCache: if cache_head_slice is None: self.k[:, local_start_index:local_end_index] = key self.v[:, local_start_index:local_end_index] = value - visible_k = self.k[:, attn_start_index:updated_local_end] - visible_v = self.v[:, attn_start_index:updated_local_end] + visible_k, visible_v = self._visible_attention_kv( + local_start_index=local_start_index, + updated_local_end=updated_local_end, + attn_start_index=attn_start_index, + recent_window_tokens=recent_window_tokens, + ) else: self.k[:, local_start_index:local_end_index, cache_head_slice, :] = key self.v[:, local_start_index:local_end_index, cache_head_slice, :] = value - visible_k = self.k[ - :, attn_start_index:updated_local_end, cache_head_slice, : - ] - visible_v = self.v[ - :, attn_start_index:updated_local_end, cache_head_slice, : - ] + visible_k, visible_v = self._visible_attention_kv( + local_start_index=local_start_index, + updated_local_end=updated_local_end, + attn_start_index=attn_start_index, + recent_window_tokens=recent_window_tokens, + cache_head_slice=cache_head_slice, + ) self._write_indices( global_end_index=updated_global_end, @@ -290,6 +301,100 @@ class CausalSelfAttentionKVCache: visible_global_end=updated_global_end, ) + def _visible_attention_kv( + self, + *, + local_start_index: int, + updated_local_end: int, + attn_start_index: int, + recent_window_tokens: int | None, + cache_head_slice: slice | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return the visible KV slice for the current attention call. + + When ``recent_window_tokens`` is ``None``, the returned token range is + the standard sliding window:: + + [attn_start_index, updated_local_end) + + When recent-window selection is enabled, ``recent_window_tokens`` must be + non-negative and the returned token ranges are:: + + sink_end = min(self.sink_tokens, updated_local_end) + recent_start = max(sink_end, local_start_index - recent_window_tokens) + [0, sink_end) + [recent_start, updated_local_end) + + Thus ``0`` keeps only sink tokens plus the current chunk. + ``cache_head_slice`` applies the same token ranges to a subset of KV + heads. + """ + if recent_window_tokens is None: + if cache_head_slice is None: + return ( + self.k[:, attn_start_index:updated_local_end], + self.v[:, attn_start_index:updated_local_end], + ) + return ( + self.k[:, attn_start_index:updated_local_end, cache_head_slice, :], + self.v[:, attn_start_index:updated_local_end, cache_head_slice, :], + ) + if recent_window_tokens < 0: + raise ValueError("recent_window_tokens must be non-negative or None") + + sink_end = min(self.sink_tokens, updated_local_end) + recent_start = max(sink_end, local_start_index - recent_window_tokens) + if recent_start <= sink_end: + if cache_head_slice is None: + return self.k[:, :updated_local_end], self.v[:, :updated_local_end] + return ( + self.k[:, :updated_local_end, cache_head_slice, :], + self.v[:, :updated_local_end, cache_head_slice, :], + ) + if sink_end <= 0: + if cache_head_slice is None: + return ( + self.k[:, recent_start:updated_local_end], + self.v[:, recent_start:updated_local_end], + ) + return ( + self.k[:, recent_start:updated_local_end, cache_head_slice, :], + self.v[:, recent_start:updated_local_end, cache_head_slice, :], + ) + + if cache_head_slice is None: + return ( + torch.cat( + [ + self.k[:, :sink_end], + self.k[:, recent_start:updated_local_end], + ], + dim=1, + ), + torch.cat( + [ + self.v[:, :sink_end], + self.v[:, recent_start:updated_local_end], + ], + dim=1, + ), + ) + return ( + torch.cat( + [ + self.k[:, :sink_end, cache_head_slice, :], + self.k[:, recent_start:updated_local_end, cache_head_slice, :], + ], + dim=1, + ), + torch.cat( + [ + self.v[:, :sink_end, cache_head_slice, :], + self.v[:, recent_start:updated_local_end, cache_head_slice, :], + ], + dim=1, + ), + ) + @dataclass(slots=True) class CrossAttentionKVCache: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py b/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py index b455e8fdb..849438704 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py @@ -79,6 +79,13 @@ from sglang.multimodal_gen.runtime.models.dits.wanvideo import ( WanTransformer3DModel, ) from sglang.multimodal_gen.runtime.models.utils import _use_aiter +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world.constants import ( + LINGBOT_C2WS_PLUCKER_EMB_CACHE, + LINGBOT_CAM_CONDITIONER_CACHE, + LINGBOT_ROPE_CACHE, + LINGBOT_SEQUENCE_SHARD_ROPE_CACHE, + LINGBOT_TIME_EMBEDDINGS_CACHE, +) from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, @@ -299,6 +306,11 @@ class LingBotWorldCausalSelfAttention(CausalWanSelfAttention): value=v, current_chunk_start=current_start, cache_head_start=cache_head_start, + recent_window_tokens=( + None + if update_cache_only + else getattr(forward_batch, "realtime_causal_kv_sample_tokens", None) + ), debug_name="LingBot KV cache", ) if update_cache_only: @@ -998,7 +1010,7 @@ class CausalLingBotWorldTransformerBlock(CausalWanTransformerBlock): return self.cam_conditioner.compute_scale_shift(c2ws_plucker_emb) cache = CausalLingBotWorldTransformer3DModel._get_request_cache( - forward_batch, "lingbot_cam_conditioner" + forward_batch, LINGBOT_CAM_CONDITIONER_CACHE ) if cache is None: return self.cam_conditioner.compute_scale_shift(c2ws_plucker_emb) @@ -1211,7 +1223,7 @@ class CausalLingBotWorldTransformer3DModel(CausalWanTransformer3DModel): ) -> torch.Tensor | None: if c2ws_plucker_emb is None: return None - cache = self._get_request_cache(forward_batch, "lingbot_c2ws_plucker_emb") + cache = self._get_request_cache(forward_batch, LINGBOT_C2WS_PLUCKER_EMB_CACHE) cache_key = ( c2ws_plucker_emb.data_ptr(), tuple(c2ws_plucker_emb.shape), @@ -1274,7 +1286,7 @@ class CausalLingBotWorldTransformer3DModel(CausalWanTransformer3DModel): start_frame: int, device: torch.device, ) -> tuple[torch.Tensor, ...]: - cache = self._get_request_cache(forward_batch, "lingbot_rope") + cache = self._get_request_cache(forward_batch, LINGBOT_ROPE_CACHE) cache_key = ( post_patch_num_frames, post_patch_height, @@ -1327,7 +1339,9 @@ class CausalLingBotWorldTransformer3DModel(CausalWanTransformer3DModel): post_patch_width: int, device: torch.device, ) -> tuple[torch.Tensor, ...]: - cache = self._get_request_cache(forward_batch, "lingbot_sequence_shard_rope") + cache = self._get_request_cache( + forward_batch, LINGBOT_SEQUENCE_SHARD_ROPE_CACHE + ) cache_key = ( local_seq_len, token_start, @@ -1393,7 +1407,7 @@ class CausalLingBotWorldTransformer3DModel(CausalWanTransformer3DModel): timestep: torch.LongTensor, forward_batch, ) -> tuple[torch.Tensor, torch.Tensor]: - cache = self._get_request_cache(forward_batch, "lingbot_time_embeddings") + cache = self._get_request_cache(forward_batch, LINGBOT_TIME_EMBEDDINGS_CACHE) current_timestep = get_forward_context().current_timestep cache_key = ( current_timestep, @@ -1426,7 +1440,7 @@ class CausalLingBotWorldTransformer3DModel(CausalWanTransformer3DModel): if not self._should_cache_cam_conditioner(forward_batch): return None - cache = self._get_request_cache(forward_batch, "lingbot_cam_conditioner") + cache = self._get_request_cache(forward_batch, LINGBOT_CAM_CONDITIONER_CACHE) if cache is None: return None diff --git a/python/sglang/multimodal_gen/runtime/pipelines/lingbot_world_causal_dmd_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/lingbot_world_causal_dmd_pipeline.py index 4ae196c8e..309539c68 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/lingbot_world_causal_dmd_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/lingbot_world_causal_dmd_pipeline.py @@ -18,7 +18,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages import ( DMDTimestepPreparationStage, ImageEncodingStage, ) -from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world import ( +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world.lingbot_world_causal_denoising import ( LingBotWorldCausalDMDDenoisingStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime import ( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index 66c2acf3d..d4d1247e6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -215,6 +215,7 @@ class Req: realtime_output_pacing: bool = False realtime_causal_sink_size: int | None = None realtime_causal_kv_cache_num_frames: int | None = None + realtime_causal_kv_sample_tokens: int | None = None # return websocket-friendly raw RGB frame bytes instead of rwa tensors return_raw_frames: bool = False diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py index 1f3458797..da9c5207f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py @@ -1094,6 +1094,7 @@ class CausalDMDDenoisingStage(DenoisingStage): device, *, sequence_shard_enabled: bool = False, + kv_cache_size: int | None = None, ) -> None: """ Initialize (but not fill) a Per-GPU KV cache aligned with the model assumptions. @@ -1102,9 +1103,10 @@ class CausalDMDDenoisingStage(DenoisingStage): sequence_shard_enabled=sequence_shard_enabled ) attention_head_dim = self.transformer.attention_head_dim - kv_cache_size = self._get_causal_kv_cache_size( - sequence_shard_enabled=sequence_shard_enabled - ) + if kv_cache_size is None: + kv_cache_size = self._get_causal_kv_cache_size( + sequence_shard_enabled=sequence_shard_enabled + ) self.causal_kv_cache = self._allocate_causal_kv_cache( batch_size=batch_size, kv_cache_size=kv_cache_size, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/__init__.py index 5e7b442f3..9029a85e5 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/__init__.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/__init__.py @@ -1,9 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 """LingBot-World-specific pipeline stages.""" -from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world.lingbot_world_causal_denoising import ( - LingBotWorldCausalDMDDenoisingStage, -) + +def __getattr__(name: str): + if name == "LingBotWorldCausalDMDDenoisingStage": + from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world.lingbot_world_causal_denoising import ( + LingBotWorldCausalDMDDenoisingStage, + ) + + return LingBotWorldCausalDMDDenoisingStage + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + __all__ = [ "LingBotWorldCausalDMDDenoisingStage", diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/constants.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/constants.py new file mode 100644 index 000000000..8d3e735cc --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/constants.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 + +LINGBOT_CAMERA_ACTIONS_CONDITION = "camera_actions" +LINGBOT_PROMPT_UPDATED_CONDITION = "lingbot_prompt_updated" + +LINGBOT_INTERACTIVE_KV_WINDOW_CACHE = "lingbot_interactive_kv_window" +LINGBOT_C2WS_PLUCKER_EMB_CACHE = "lingbot_c2ws_plucker_emb" +LINGBOT_CAM_CONDITIONER_CACHE = "lingbot_cam_conditioner" +LINGBOT_ROPE_CACHE = "lingbot_rope" +LINGBOT_SEQUENCE_SHARD_ROPE_CACHE = "lingbot_sequence_shard_rope" +LINGBOT_TIME_EMBEDDINGS_CACHE = "lingbot_time_embeddings" diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/lingbot_world_causal_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/lingbot_world_causal_denoising.py index 263166af3..44615e2c8 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/lingbot_world_causal_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/lingbot_world_causal_denoising.py @@ -3,8 +3,11 @@ """LingBot-World causal DMD denoising stage.""" +from typing import Any + import torch +from sglang.multimodal_gen import envs from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_ring_parallel_world_size, get_ulysses_parallel_world_size, @@ -14,6 +17,15 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.causal_denoising import ( CausalDMDCachePolicy, CausalDMDDenoisingStage, + CausalDMDForwardContext, + CausalDMDRealtimeCacheContext, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world.constants import ( + LINGBOT_C2WS_PLUCKER_EMB_CACHE, + LINGBOT_CAM_CONDITIONER_CACHE, + LINGBOT_CAMERA_ACTIONS_CONDITION, + LINGBOT_INTERACTIVE_KV_WINDOW_CACHE, + LINGBOT_PROMPT_UPDATED_CONDITION, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( StageValidators as V, @@ -23,6 +35,9 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) class LingBotWorldCausalDMDDenoisingStage(CausalDMDDenoisingStage): @@ -76,8 +91,11 @@ class LingBotWorldCausalDMDDenoisingStage(CausalDMDDenoisingStage): def _causal_kv_cache_kwargs( self, policy: CausalDMDCachePolicy, - ) -> dict[str, bool]: - return {"sequence_shard_enabled": policy.sequence_shard_enabled} + ) -> dict[str, Any]: + return { + "sequence_shard_enabled": policy.sequence_shard_enabled, + "kv_cache_size": policy.expected_cache_tokens, + } def _use_causal_cache_int_indices( self, @@ -86,6 +104,316 @@ class LingBotWorldCausalDMDDenoisingStage(CausalDMDDenoisingStage): ) -> bool: return True + @staticmethod + def _chunk_has_camera_motion(actions) -> bool: + if not actions: + return False + for frame_actions in actions: + if frame_actions: + return True + return False + + def _uses_interactive_kv_window( + self, + batch: Req, + server_args: ServerArgs, + ) -> bool: + if not self._interactive_kv_window_enabled(server_args): + return False + condition_inputs = getattr(batch, "condition_inputs", None) or {} + return LINGBOT_CAMERA_ACTIONS_CONDITION in condition_inputs + + @staticmethod + def _interactive_kv_window_enabled(server_args: ServerArgs) -> bool: + config_enabled = bool( + getattr( + server_args.pipeline_config, + "interactive_kv_window_enable", + False, + ) + ) + return config_enabled or envs.SGLANG_LINGBOT_ENABLE_INTERACTIVE_KV_WINDOW + + def _apply_causal_cache_overrides( + self, + batch: Req, + server_args: ServerArgs, + ) -> None: + self._reset_causal_cache_config_defaults() + super()._apply_causal_cache_overrides(batch, server_args) + self._sync_interactive_kv_cache_window(server_args) + + def _reset_causal_cache_config_defaults(self) -> None: + arch_config = getattr( + getattr(self.transformer, "config", None), "arch_config", None + ) + if arch_config is None: + return + if hasattr(arch_config, "sink_size"): + self.sink_size = int(arch_config.sink_size) + if hasattr(arch_config, "sliding_window_num_frames"): + self.sliding_window_num_frames = int(arch_config.sliding_window_num_frames) + + def _sync_interactive_kv_cache_window(self, server_args: ServerArgs) -> None: + if not self._interactive_kv_window_enabled(server_args): + return + if self.local_attn_size != -1: + return + self.sliding_window_num_frames = ( + self._effective_interactive_kv_cache_num_frames(server_args) + ) + + def _effective_interactive_kv_cache_num_frames( + self, + server_args: ServerArgs, + ) -> int: + cache_window = int(self.sliding_window_num_frames) + if self.local_attn_size != -1: + return cache_window + + moving_window = self._moving_kv_sample_num_frames(server_args) or 0 + still_window = self._still_kv_sample_num_frames(server_args) or 0 + return max( + cache_window, + int(self.sink_size) + + max(moving_window, still_window) + + int(self.num_frames_per_block), + ) + + def _build_realtime_causal_cache_policy( + self, + batch: Req, + server_args: ServerArgs, + ) -> CausalDMDCachePolicy: + policy = super()._build_realtime_causal_cache_policy(batch, server_args) + if self._interactive_kv_window_enabled(server_args): + policy.expected_cache_tokens = ( + self._effective_interactive_kv_cache_num_frames(server_args) + * self.num_token_per_frame + ) + return policy + + @staticmethod + def _should_reset_lingbot_crossattn_cache(batch: Req) -> bool: + condition_inputs = getattr(batch, "condition_inputs", None) or {} + return bool(condition_inputs.get(LINGBOT_PROMPT_UPDATED_CONDITION)) + + def _sync_lingbot_crossattn_cache( + self, + batch: Req, + cache_ctx: CausalDMDRealtimeCacheContext, + ) -> None: + if self._should_reset_lingbot_crossattn_cache(batch): + self._reset_crossattn_cache(cache_ctx.crossattn_cache) + + def _prepare_realtime_causal_caches( + self, + batch: Req, + server_args: ServerArgs, + ctx: CausalDMDForwardContext, + ) -> CausalDMDRealtimeCacheContext: + cache_ctx = super()._prepare_realtime_causal_caches(batch, server_args, ctx) + self._sync_lingbot_crossattn_cache(batch, cache_ctx) + return cache_ctx + + def _base_kv_sample_num_frames(self) -> int | None: + sample_frames = ( + int(self.sliding_window_num_frames) + - int(self.sink_size) + - int(self.num_frames_per_block) + ) + return sample_frames if sample_frames > 0 else None + + @staticmethod + def _optional_non_negative_int(value: Any) -> int | None: + if value is None: + return None + return max(0, int(value)) + + def _moving_kv_sample_num_frames( + self, + server_args: ServerArgs, + ) -> int | None: + return self._optional_non_negative_int( + getattr( + server_args.pipeline_config, + "interactive_kv_moving_window", + None, + ) + ) + + def _still_kv_sample_num_frames( + self, + server_args: ServerArgs, + ) -> int | None: + return self._optional_non_negative_int( + getattr( + server_args.pipeline_config, + "interactive_kv_still_window", + 3, + ) + ) + + def _get_interactive_kv_sample_num_frames( + self, + cache_state, + batch: Req, + server_args: ServerArgs, + ) -> int | None: + pipeline_config = server_args.pipeline_config + if not self._interactive_kv_window_enabled(server_args): + return None + if not self._uses_interactive_kv_window(batch, server_args): + return self._base_kv_sample_num_frames() + + dynamic_state = cache_state.runtime_cache.setdefault( + LINGBOT_INTERACTIVE_KV_WINDOW_CACHE, + { + "consecutive_still_chunks": 0, + "sample_num_frames": None, + }, + ) + if cache_state.chunk_idx == 0: + dynamic_state["consecutive_still_chunks"] = 0 + dynamic_state["sample_num_frames"] = None + + moving_window = self._moving_kv_sample_num_frames(server_args) + if moving_window is None: + return None + still_window = self._still_kv_sample_num_frames(server_args) + still_chunks_threshold = max( + 1, int(getattr(pipeline_config, "interactive_kv_still_chunks", 2)) + ) + if dynamic_state["sample_num_frames"] is None: + dynamic_state["sample_num_frames"] = moving_window + + condition_inputs = getattr(batch, "condition_inputs", None) or {} + if self._chunk_has_camera_motion( + condition_inputs.get(LINGBOT_CAMERA_ACTIONS_CONDITION) + ): + dynamic_state["consecutive_still_chunks"] = 0 + dynamic_state["sample_num_frames"] = moving_window + else: + dynamic_state["consecutive_still_chunks"] += 1 + if ( + still_window is not None + and dynamic_state["consecutive_still_chunks"] >= still_chunks_threshold + ): + dynamic_state["sample_num_frames"] = still_window + + return int(dynamic_state["sample_num_frames"]) + + def _log_lingbot_kv_window( + self, + cache_state, + batch: Req, + server_args: ServerArgs, + *, + sample_frames: int | None, + ) -> None: + if not self._interactive_kv_window_enabled(server_args): + return + + mode = "base" + still_chunks = None + if self._uses_interactive_kv_window(batch, server_args): + dynamic_state = cache_state.runtime_cache.get( + LINGBOT_INTERACTIVE_KV_WINDOW_CACHE, {} + ) + still_chunks = dynamic_state.get("consecutive_still_chunks") + condition_inputs = getattr(batch, "condition_inputs", None) or {} + if self._chunk_has_camera_motion( + condition_inputs.get(LINGBOT_CAMERA_ACTIONS_CONDITION) + ): + mode = "moving" + else: + still_window = self._still_kv_sample_num_frames(server_args) + still_chunks_threshold = max( + 1, + int( + getattr( + server_args.pipeline_config, + "interactive_kv_still_chunks", + 2, + ) + ), + ) + if ( + still_window is not None + and sample_frames == still_window + and still_chunks is not None + and still_chunks >= still_chunks_threshold + ): + mode = "still" + else: + mode = "moving" + + window_frames = ( + int(self.sliding_window_num_frames) + if sample_frames is None + else int(self.sink_size) + + int(sample_frames) + + int(self.num_frames_per_block) + ) + sample_tokens = ( + None + if sample_frames is None + else int(sample_frames) * int(self.num_token_per_frame) + ) + logger.info( + "LingBot interactive KV window: session_id=%s request_id=%s " + "chunk_idx=%s mode=%s window_frames=%s sample_frames=%s " + "cache_frames=%s sink_frames=%s current_frames=%s sample_tokens=%s " + "cache_tokens=%s still_chunks=%s", + getattr(batch, "realtime_session_id", None), + getattr(batch, "request_id", None), + getattr(batch, "block_idx", None), + mode, + window_frames, + sample_frames, + int(self.sliding_window_num_frames), + int(self.sink_size), + int(self.num_frames_per_block), + sample_tokens, + int(self.sliding_window_num_frames) * int(self.num_token_per_frame), + still_chunks, + ) + + def _set_lingbot_kv_sample_tokens( + self, + cache_state, + batch: Req, + server_args: ServerArgs, + ) -> int | None: + self._sync_interactive_kv_cache_window(server_args) + sample_frames = self._get_interactive_kv_sample_num_frames( + cache_state, + batch, + server_args, + ) + sample_tokens = ( + None + if sample_frames is None + else int(sample_frames) * self.num_token_per_frame + ) + self._log_lingbot_kv_window( + cache_state, + batch, + server_args, + sample_frames=sample_frames, + ) + previous = getattr(batch, "realtime_causal_kv_sample_tokens", None) + batch.realtime_causal_kv_sample_tokens = sample_tokens + return previous + + @staticmethod + def _clear_lingbot_dynamic_condition_cache(cache_state) -> None: + runtime_cache = getattr(cache_state, "runtime_cache", None) + if runtime_cache is None: + return + runtime_cache.pop(LINGBOT_C2WS_PLUCKER_EMB_CACHE, None) + runtime_cache.pop(LINGBOT_CAM_CONDITIONER_CACHE, None) + def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: result = VerificationResult() result.add_check( @@ -97,6 +425,35 @@ class LingBotWorldCausalDMDDenoisingStage(CausalDMDDenoisingStage): result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) return result + def _denoise_realtime_causal_chunk( + self, + batch: Req, + server_args: ServerArgs, + *, + ctx, + cache_ctx, + chunk_latents: torch.Tensor, + prepare_model_input, + prepare_context_input, + ) -> torch.Tensor: + previous_sample_tokens = self._set_lingbot_kv_sample_tokens( + cache_ctx.cache_state, + batch, + server_args, + ) + try: + return super()._denoise_realtime_causal_chunk( + batch, + server_args, + ctx=ctx, + cache_ctx=cache_ctx, + chunk_latents=chunk_latents, + prepare_model_input=prepare_model_input, + prepare_context_input=prepare_context_input, + ) + finally: + batch.realtime_causal_kv_sample_tokens = previous_sample_tokens + def _get_causal_dmd_latents(self, batch: Req) -> torch.Tensor: latents = batch.latents assert latents is not None, ( @@ -294,16 +651,18 @@ class LingBotWorldCausalDMDDenoisingStage(CausalDMDDenoisingStage): device=ctx.device, ) - current_latents = self._denoise_realtime_causal_chunk( - batch, - server_args, - ctx=ctx, - cache_ctx=cache_ctx, - chunk_latents=current_latents, - prepare_model_input=prepare_model_input, - prepare_context_input=prepare_model_input, - ) - cache_ctx.cache_state.runtime_cache.pop("lingbot_cam_conditioner", None) + try: + current_latents = self._denoise_realtime_causal_chunk( + batch, + server_args, + ctx=ctx, + cache_ctx=cache_ctx, + chunk_latents=current_latents, + prepare_model_input=prepare_model_input, + prepare_context_input=prepare_model_input, + ) + finally: + self._clear_lingbot_dynamic_condition_cache(cache_ctx.cache_state) # Advance cumulative frame position self._advance_realtime_causal_cache(cache_ctx, num_frames=ctx.num_frames) diff --git a/python/sglang/multimodal_gen/test/unit/realtime/test_lingbot_causal_denoising.py b/python/sglang/multimodal_gen/test/unit/realtime/test_lingbot_causal_denoising.py index cdc745914..308b209c1 100644 --- a/python/sglang/multimodal_gen/test/unit/realtime/test_lingbot_causal_denoising.py +++ b/python/sglang/multimodal_gen/test/unit/realtime/test_lingbot_causal_denoising.py @@ -4,7 +4,11 @@ from types import SimpleNamespace import torch +from sglang.multimodal_gen.configs.pipeline_configs.lingbot_world import ( + LingBotWorldCausalDMDConfig, +) from sglang.multimodal_gen.runtime.layers.kvcache.causal_attention_cache import ( + CausalSelfAttentionKVCache, CrossAttentionKVCache, ) from sglang.multimodal_gen.runtime.models.dits import ( @@ -15,9 +19,21 @@ from sglang.multimodal_gen.runtime.models.dits.lingbot_world import ( CausalLingBotWorldTransformerBlock, LingBotWorldCamConditioner, ) -from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world import ( +from sglang.multimodal_gen.runtime.pipelines_core.stages.causal_denoising import ( + CausalDMDCachePolicy, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world.constants import ( + LINGBOT_C2WS_PLUCKER_EMB_CACHE, + LINGBOT_CAM_CONDITIONER_CACHE, + LINGBOT_PROMPT_UPDATED_CONDITION, + LINGBOT_ROPE_CACHE, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world.lingbot_world_causal_denoising import ( LingBotWorldCausalDMDDenoisingStage, ) +from sglang.multimodal_gen.runtime.realtime.states import RealtimeCausalDiTState + +LINGBOT_INTERACTIVE_KV_WINDOW_ENV = "SGLANG_LINGBOT_ENABLE_INTERACTIVE_KV_WINDOW" def test_lingbot_denoising_stage_does_not_own_realtime_cache_refs(): @@ -181,6 +197,408 @@ def test_lingbot_realtime_attention_cache_rolls_with_sink_window(): ] +def test_lingbot_realtime_attention_cache_samples_sink_and_recent_window(): + cache = CausalSelfAttentionKVCache( + k=torch.zeros(1, 8, 1, 1), + v=torch.zeros(1, 8, 1, 1), + global_end_index=torch.zeros(1, dtype=torch.long), + local_end_index=torch.zeros(1, dtype=torch.long), + global_end_index_int=0, + local_end_index_int=0, + cache_size=8, + sink_tokens=2, + attention_window_size=8, + ) + + cache.update_and_get_attention_kv( + key=torch.ones(1, 3, 1, 1), + value=torch.ones(1, 3, 1, 1), + current_chunk_start=0, + ) + cache.update_and_get_attention_kv( + key=torch.full((1, 3, 1, 1), 2.0), + value=torch.full((1, 3, 1, 1), 2.0), + current_chunk_start=3, + ) + sampled_view = cache.update_and_get_attention_kv( + key=torch.full((1, 3, 1, 1), 3.0), + value=torch.full((1, 3, 1, 1), 3.0), + current_chunk_start=6, + recent_window_tokens=1, + ) + + assert sampled_view.k.flatten().tolist() == [1.0, 1.0, 2.0, 3.0, 3.0, 3.0] + assert sampled_view.v.flatten().tolist() == [1.0, 1.0, 2.0, 3.0, 3.0, 3.0] + + current_only_view = cache.update_and_get_attention_kv( + key=torch.full((1, 3, 1, 1), 4.0), + value=torch.full((1, 3, 1, 1), 4.0), + current_chunk_start=6, + recent_window_tokens=0, + ) + + assert current_only_view.k.flatten().tolist() == [1.0, 1.0, 4.0, 4.0, 4.0] + assert current_only_view.v.flatten().tolist() == [1.0, 1.0, 4.0, 4.0, 4.0] + + +def test_lingbot_interactive_kv_window_config_default_disabled(): + field = LingBotWorldCausalDMDConfig.__dataclass_fields__[ + "interactive_kv_window_enable" + ] + + assert field.default is False + + +def test_lingbot_lazy_vae_encode_black_frames_env(monkeypatch): + config = LingBotWorldCausalDMDConfig() + temporal_ratio = int(config.vae_config.arch_config.temporal_compression_ratio) + image = torch.zeros(1, 3, temporal_ratio * 2 + 1, 2, 2) + + monkeypatch.delenv("SGLANG_LINGBOT_LAZY_VAE_ENCODE_BLACK_FRAMES", raising=False) + assert config.preprocess_vae_encode(image, vae=None) is image + + monkeypatch.setenv("SGLANG_LINGBOT_LAZY_VAE_ENCODE_BLACK_FRAMES", "0") + assert config.preprocess_vae_encode(image, vae=None) is image + + monkeypatch.setenv( + "SGLANG_LINGBOT_LAZY_VAE_ENCODE_BLACK_FRAMES", str(temporal_ratio) + ) + encoded = config.preprocess_vae_encode(image, vae=None) + assert encoded.shape[2] == temporal_ratio + 1 + + +def test_lingbot_interactive_kv_window_samples_base_moving_and_still(monkeypatch): + monkeypatch.delenv(LINGBOT_INTERACTIVE_KV_WINDOW_ENV, raising=False) + stage = LingBotWorldCausalDMDDenoisingStage.__new__( + LingBotWorldCausalDMDDenoisingStage + ) + stage.local_attn_size = -1 + stage.sink_size = 9 + stage.num_token_per_frame = 10 + stage.num_frames_per_block = 3 + stage.sliding_window_num_frames = 18 + stage.transformer = SimpleNamespace(num_attention_heads=1) + server_args = SimpleNamespace( + pipeline_config=SimpleNamespace( + interactive_kv_window_enable=True, + interactive_kv_moving_window=12, + interactive_kv_still_window=3, + interactive_kv_still_chunks=2, + ) + ) + cache_state = RealtimeCausalDiTState() + + batch = SimpleNamespace(condition_inputs={}) + previous = stage._set_lingbot_kv_sample_tokens(cache_state, batch, server_args) + assert previous is None + assert stage.sliding_window_num_frames == 24 + assert batch.realtime_causal_kv_sample_tokens == 120 + policy = stage._build_realtime_causal_cache_policy(batch, server_args) + assert policy.expected_cache_tokens == 240 + + batch.condition_inputs = {"camera_actions": [["w"], [], []]} + cache_state.chunk_idx = 0 + stage._set_lingbot_kv_sample_tokens(cache_state, batch, server_args) + assert batch.realtime_causal_kv_sample_tokens == 120 + + batch.condition_inputs = {"camera_actions": [[], [], []]} + cache_state.chunk_idx = 1 + stage._set_lingbot_kv_sample_tokens(cache_state, batch, server_args) + assert batch.realtime_causal_kv_sample_tokens == 120 + + cache_state.chunk_idx = 2 + stage._set_lingbot_kv_sample_tokens(cache_state, batch, server_args) + assert batch.realtime_causal_kv_sample_tokens == 30 + + +def test_lingbot_interactive_kv_window_none_disables_moving_window(monkeypatch): + monkeypatch.delenv(LINGBOT_INTERACTIVE_KV_WINDOW_ENV, raising=False) + stage = LingBotWorldCausalDMDDenoisingStage.__new__( + LingBotWorldCausalDMDDenoisingStage + ) + stage.local_attn_size = -1 + stage.sink_size = 9 + stage.num_token_per_frame = 10 + stage.num_frames_per_block = 3 + stage.sliding_window_num_frames = 18 + stage.transformer = SimpleNamespace(num_attention_heads=1) + server_args = SimpleNamespace( + pipeline_config=SimpleNamespace( + realtime_causal_sink_size=9, + realtime_causal_kv_cache_num_frames=18, + interactive_kv_window_enable=True, + interactive_kv_moving_window=None, + interactive_kv_still_window=3, + interactive_kv_still_chunks=2, + ) + ) + cache_state = RealtimeCausalDiTState() + batch = SimpleNamespace(condition_inputs={"camera_actions": [["w"], [], []]}) + + stage._set_lingbot_kv_sample_tokens(cache_state, batch, server_args) + assert batch.realtime_causal_kv_sample_tokens is None + policy = stage._build_realtime_causal_cache_policy(batch, server_args) + assert policy.expected_cache_tokens == 180 + + +def test_lingbot_interactive_kv_window_zero_is_valid_moving_window(monkeypatch): + monkeypatch.delenv(LINGBOT_INTERACTIVE_KV_WINDOW_ENV, raising=False) + stage = LingBotWorldCausalDMDDenoisingStage.__new__( + LingBotWorldCausalDMDDenoisingStage + ) + stage.local_attn_size = -1 + stage.sink_size = 9 + stage.num_token_per_frame = 10 + stage.num_frames_per_block = 3 + stage.sliding_window_num_frames = 18 + stage.transformer = SimpleNamespace(num_attention_heads=1) + server_args = SimpleNamespace( + pipeline_config=SimpleNamespace( + realtime_causal_sink_size=9, + realtime_causal_kv_cache_num_frames=18, + interactive_kv_window_enable=True, + interactive_kv_moving_window=0, + interactive_kv_still_window=None, + interactive_kv_still_chunks=2, + ) + ) + cache_state = RealtimeCausalDiTState() + batch = SimpleNamespace(condition_inputs={"camera_actions": [["w"], [], []]}) + + stage._set_lingbot_kv_sample_tokens(cache_state, batch, server_args) + assert batch.realtime_causal_kv_sample_tokens == 0 + policy = stage._build_realtime_causal_cache_policy(batch, server_args) + assert policy.expected_cache_tokens == 180 + + +def test_lingbot_interactive_kv_window_updates_total_window_for_moving_default( + monkeypatch, +): + monkeypatch.delenv(LINGBOT_INTERACTIVE_KV_WINDOW_ENV, raising=False) + stage = LingBotWorldCausalDMDDenoisingStage.__new__( + LingBotWorldCausalDMDDenoisingStage + ) + stage.local_attn_size = -1 + stage.sink_size = 9 + stage.num_token_per_frame = 10 + stage.num_frames_per_block = 3 + stage.sliding_window_num_frames = 18 + stage.transformer = SimpleNamespace(num_attention_heads=1) + server_args = SimpleNamespace( + pipeline_config=SimpleNamespace( + realtime_causal_sink_size=9, + realtime_causal_kv_cache_num_frames=18, + interactive_kv_window_enable=True, + interactive_kv_moving_window=12, + interactive_kv_still_window=3, + interactive_kv_still_chunks=2, + ) + ) + batch = SimpleNamespace(condition_inputs={"camera_actions": [["w"], [], []]}) + + policy = stage._build_realtime_causal_cache_policy(batch, server_args) + + assert stage.sliding_window_num_frames == 24 + assert policy.expected_cache_tokens == 240 + + +def test_lingbot_interactive_kv_window_resets_stage_window_between_requests( + monkeypatch, +): + monkeypatch.delenv(LINGBOT_INTERACTIVE_KV_WINDOW_ENV, raising=False) + stage = LingBotWorldCausalDMDDenoisingStage.__new__( + LingBotWorldCausalDMDDenoisingStage + ) + stage.local_attn_size = -1 + stage.sink_size = 9 + stage.num_token_per_frame = 10 + stage.num_frames_per_block = 3 + stage.sliding_window_num_frames = 18 + stage.transformer = SimpleNamespace( + num_attention_heads=1, + config=SimpleNamespace( + arch_config=SimpleNamespace( + sink_size=9, + sliding_window_num_frames=18, + ) + ), + ) + dynamic_args = SimpleNamespace( + pipeline_config=SimpleNamespace( + realtime_causal_sink_size=9, + realtime_causal_kv_cache_num_frames=18, + interactive_kv_window_enable=True, + interactive_kv_moving_window=12, + interactive_kv_still_window=3, + interactive_kv_still_chunks=2, + ) + ) + disabled_args = SimpleNamespace( + pipeline_config=SimpleNamespace( + interactive_kv_window_enable=False, + interactive_kv_moving_window=12, + interactive_kv_still_window=3, + interactive_kv_still_chunks=2, + ) + ) + + dynamic_batch = SimpleNamespace(condition_inputs={"camera_actions": [["w"]]}) + dynamic_policy = stage._build_realtime_causal_cache_policy( + dynamic_batch, dynamic_args + ) + assert stage.sliding_window_num_frames == 24 + assert dynamic_policy.expected_cache_tokens == 240 + + disabled_batch = SimpleNamespace(condition_inputs={}) + disabled_policy = stage._build_realtime_causal_cache_policy( + disabled_batch, disabled_args + ) + assert stage.sliding_window_num_frames == 18 + assert disabled_policy.expected_cache_tokens == 180 + + +def test_lingbot_interactive_kv_window_default_disabled(monkeypatch): + monkeypatch.delenv(LINGBOT_INTERACTIVE_KV_WINDOW_ENV, raising=False) + stage = LingBotWorldCausalDMDDenoisingStage.__new__( + LingBotWorldCausalDMDDenoisingStage + ) + stage.local_attn_size = -1 + stage.sink_size = 9 + stage.num_token_per_frame = 10 + stage.num_frames_per_block = 3 + stage.sliding_window_num_frames = 18 + stage.transformer = SimpleNamespace(num_attention_heads=1) + server_args = SimpleNamespace( + pipeline_config=SimpleNamespace( + realtime_causal_sink_size=9, + realtime_causal_kv_cache_num_frames=18, + interactive_kv_window_enable=False, + interactive_kv_moving_window=12, + interactive_kv_still_window=3, + interactive_kv_still_chunks=2, + ) + ) + cache_state = RealtimeCausalDiTState() + batch = SimpleNamespace(condition_inputs={"camera_actions": [["w"], [], []]}) + + stage._set_lingbot_kv_sample_tokens(cache_state, batch, server_args) + policy = stage._build_realtime_causal_cache_policy(batch, server_args) + + assert stage.sliding_window_num_frames == 18 + assert batch.realtime_causal_kv_sample_tokens is None + assert policy.expected_cache_tokens == 180 + + +def test_lingbot_interactive_kv_window_env_can_enable_default(monkeypatch): + monkeypatch.setenv(LINGBOT_INTERACTIVE_KV_WINDOW_ENV, "1") + stage = LingBotWorldCausalDMDDenoisingStage.__new__( + LingBotWorldCausalDMDDenoisingStage + ) + stage.local_attn_size = -1 + stage.sink_size = 9 + stage.num_token_per_frame = 10 + stage.num_frames_per_block = 3 + stage.sliding_window_num_frames = 18 + stage.transformer = SimpleNamespace(num_attention_heads=1) + server_args = SimpleNamespace( + pipeline_config=SimpleNamespace( + realtime_causal_sink_size=9, + realtime_causal_kv_cache_num_frames=18, + interactive_kv_window_enable=False, + interactive_kv_moving_window=12, + interactive_kv_still_window=3, + interactive_kv_still_chunks=2, + ) + ) + cache_state = RealtimeCausalDiTState() + batch = SimpleNamespace(condition_inputs={"camera_actions": [["w"], [], []]}) + + stage._set_lingbot_kv_sample_tokens(cache_state, batch, server_args) + policy = stage._build_realtime_causal_cache_policy(batch, server_args) + + assert stage.sliding_window_num_frames == 24 + assert batch.realtime_causal_kv_sample_tokens == 120 + assert policy.expected_cache_tokens == 240 + + +def test_lingbot_interactive_kv_window_allocates_expected_cache_size(): + stage = LingBotWorldCausalDMDDenoisingStage.__new__( + LingBotWorldCausalDMDDenoisingStage + ) + stage.num_transformer_blocks = 1 + stage.local_attn_size = -1 + stage.sink_size = 9 + stage.num_token_per_frame = 10 + stage.num_frames_per_block = 3 + stage.sliding_window_num_frames = 18 + stage.transformer = SimpleNamespace(num_attention_heads=1, attention_head_dim=1) + policy = CausalDMDCachePolicy( + sequence_shard_enabled=False, + num_attention_heads=1, + expected_cache_tokens=240, + expected_sink_tokens=90, + kv_cache_kwargs={}, + ) + + stage._initialize_kv_cache( + batch_size=1, + dtype=torch.float32, + device=torch.device("cpu"), + **stage._causal_kv_cache_kwargs(policy), + ) + + assert stage.causal_kv_cache is not None + cache = stage.causal_kv_cache[0] + assert cache.cache_size == 240 + assert cache.k.shape[1] == 240 + assert cache.sink_tokens == 90 + assert cache.attention_window_size == 240 + + +def test_lingbot_dynamic_condition_cache_clear_removes_chunk_entries(): + cache_state = SimpleNamespace( + runtime_cache={ + LINGBOT_C2WS_PLUCKER_EMB_CACHE: object(), + LINGBOT_CAM_CONDITIONER_CACHE: object(), + LINGBOT_ROPE_CACHE: object(), + } + ) + + LingBotWorldCausalDMDDenoisingStage._clear_lingbot_dynamic_condition_cache( + cache_state + ) + + assert LINGBOT_C2WS_PLUCKER_EMB_CACHE not in cache_state.runtime_cache + assert LINGBOT_CAM_CONDITIONER_CACHE not in cache_state.runtime_cache + assert LINGBOT_ROPE_CACHE in cache_state.runtime_cache + + +def test_lingbot_crossattn_cache_resets_on_prompt_event(): + stage = LingBotWorldCausalDMDDenoisingStage.__new__( + LingBotWorldCausalDMDDenoisingStage + ) + crossattn_cache = [ + CrossAttentionKVCache( + k=torch.empty(1), + v=torch.empty(1), + ) + ] + crossattn_cache[0].store(torch.ones(1), torch.ones(1)) + cache_ctx = SimpleNamespace( + cache_state=RealtimeCausalDiTState(), + crossattn_cache=crossattn_cache, + ) + batch = SimpleNamespace(condition_inputs={}) + + stage._sync_lingbot_crossattn_cache(batch, cache_ctx) + assert crossattn_cache[0].is_init + + batch.condition_inputs = {LINGBOT_PROMPT_UPDATED_CONDITION: True} + stage._sync_lingbot_crossattn_cache(batch, cache_ctx) + + assert not crossattn_cache[0].is_init + + def test_lingbot_i2v_model_input_writer_reuses_buffer(): latents = torch.ones(1, 16, 3, 2, 2) condition = torch.full((1, 20, 3, 2, 2), 2.0) @@ -255,7 +673,7 @@ def test_lingbot_cam_conditioner_cache_reuses_source_tensor(monkeypatch): assert first is second assert third is not first assert block.cam_conditioner.calls == 2 - cache = forward_batch.extra["lingbot_cam_conditioner"] + cache = forward_batch.extra[LINGBOT_CAM_CONDITIONER_CACHE] assert cache["source_key"][0] == next_source.data_ptr() assert len(cache["entries"]) == 1 @@ -287,7 +705,7 @@ def test_lingbot_cam_conditioner_cache_skips_non_sequence_shard(monkeypatch): assert first is not second assert first[0] is not second[0] assert block.cam_conditioner.calls == 2 - assert "lingbot_cam_conditioner" not in forward_batch.extra + assert LINGBOT_CAM_CONDITIONER_CACHE not in forward_batch.extra def test_lingbot_cam_conditioner_cache_skips_single_ulysses_world(monkeypatch): @@ -319,7 +737,7 @@ def test_lingbot_cam_conditioner_cache_skips_single_ulysses_world(monkeypatch): assert first is not second assert block.cam_conditioner.calls == 2 - assert "lingbot_cam_conditioner" not in forward_batch.extra + assert LINGBOT_CAM_CONDITIONER_CACHE not in forward_batch.extra def test_lingbot_cam_conditioner_cache_reuses_context_update(monkeypatch): @@ -351,7 +769,7 @@ def test_lingbot_cam_conditioner_cache_reuses_context_update(monkeypatch): assert first is second assert block.cam_conditioner.calls == 1 - assert "lingbot_cam_conditioner" in forward_batch.extra + assert LINGBOT_CAM_CONDITIONER_CACHE in forward_batch.extra def test_lingbot_model_prepares_cam_conditioner_scale_shifts(monkeypatch): diff --git a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py index 29216d062..4cecc9755 100644 --- a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py +++ b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py @@ -6,6 +6,7 @@ from types import SimpleNamespace import msgspec.msgpack import numpy as np +import pytest import torch from sglang.multimodal_gen.configs.pipeline_configs.lingbot_world import ( @@ -36,7 +37,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.registry import ( get_realtime_model_adapter, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch -from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world import ( +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world.lingbot_world_causal_denoising import ( LingBotWorldCausalDMDDenoisingStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.base import ( @@ -296,6 +297,104 @@ def test_lingbot_realtime_adapter_ingests_generic_events(): assert state.latest_sampled_event_id == 8 +def test_lingbot_realtime_adapter_ingests_composite_input_event(): + adapter = lingbot_realtime.LingBotWorldRealtimeAdapter() + session = GenerateSession() + session.set_adapter(adapter) + session.set_request( + RealtimeVideoGenerationsRequest( + type="init", + prompt="walk forward", + ) + ) + + composite_event = RealtimeEvent( + type="event", + kind="composite_input", + payload={ + "input_types": ["prompt", "camera_actions"], + "prompt": "turn left", + "camera_actions": [["w"], ["d"]], + }, + event_id=9, + ) + + event_log = adapter.ingest_event(session, composite_event) + + assert "kind=composite_input" in event_log + chunk_inputs = adapter.sample_chunk_inputs( + session, + server_args=SimpleNamespace(), + chunk=SimpleNamespace(index=1), + chunk_size=3, + ) + assert chunk_inputs.prompt == "turn left" + assert chunk_inputs.condition_inputs[ + lingbot_realtime.LINGBOT_PROMPT_UPDATED_CONDITION + ] + assert chunk_inputs.condition_inputs[ + lingbot_realtime.LINGBOT_CAMERA_ACTIONS_CONDITION + ] == [["w"], ["d"], []] + assert adapter.get_realtime_event_id(session) == 9 + + +def test_lingbot_realtime_adapter_rejects_composite_input_atomically(): + adapter = lingbot_realtime.LingBotWorldRealtimeAdapter() + session = GenerateSession() + session.set_adapter(adapter) + session.set_request( + RealtimeVideoGenerationsRequest( + type="init", + prompt="walk forward", + ) + ) + + composite_event = RealtimeEvent( + type="event", + kind="composite_input", + payload={ + "input_types": ["prompt", "camera_actions"], + "prompt": "turn left", + "camera_actions": ["w"], + }, + event_id=10, + ) + + with pytest.raises(ValueError, match="camera_actions"): + adapter.ingest_event(session, composite_event) + + state = adapter._state(session) + assert not state.has_prompt() + assert state.sample_camera_actions(3) is None + + +def test_lingbot_realtime_prompt_event_marks_crossattn_reset(): + adapter = lingbot_realtime.LingBotWorldRealtimeAdapter() + session = GenerateSession() + session.set_adapter(adapter) + session.set_request( + RealtimeVideoGenerationsRequest( + type="init", + prompt="walk forward", + ) + ) + state = adapter._state(session) + state.receive_prompt("turn left", event_id=8) + + chunk_inputs = adapter.sample_chunk_inputs( + session, + server_args=SimpleNamespace(), + chunk=SimpleNamespace(index=1), + chunk_size=3, + ) + + assert chunk_inputs.prompt == "turn left" + assert session.request.prompt == "turn left" + assert chunk_inputs.condition_inputs[ + lingbot_realtime.LINGBOT_PROMPT_UPDATED_CONDITION + ] + + def test_lingbot_realtime_adapter_ingests_state_camera_events(): adapter = lingbot_realtime.LingBotWorldRealtimeAdapter() session = GenerateSession() @@ -672,6 +771,7 @@ def test_lingbot_realtime_condition_horizon_repeats_blank_tail_chunk(): latent_condition = torch.ones(1, latent_channels, latent_frames, 2, 2) batch = SimpleNamespace( height=2 * spatial_ratio, + num_frames=num_frames, width=2 * spatial_ratio, ) condition_full = config.postprocess_image_latent(latent_condition, batch)