minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] feat: add LingBot realtime prompt, KV window, and lazy VAE controls (#30040)

Browse Source
This commit is contained in:
HuangJi
2026-07-05 11:11:54 +08:00
committed by GitHub
parent fbe3110866
commit 5e6f49c986
13 changed files with 1286 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/pipeline_configs/lingbot_world.py
+64 -3
View File
@@ -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...])
python/sglang/multimodal_gen/envs.py
+8
View File
@@ -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:
python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/lingbot_world_realtime_adapter.py
+158 -15
View File
@@ -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(
python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py
+113 -8
View File
@@ -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:
python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py
+20 -6
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines/lingbot_world_causal_dmd_pipeline.py
+1 -1
View File
@@ -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 (
python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py
+1
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py
+5 -3
View File
@@ -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,
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/__init__.py
+10 -3
View File
@@ -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",
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/constants.py
+11
View File
@@ -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"
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/lingbot_world_causal_denoising.py
+371 -12
View File
@@ -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)
python/sglang/multimodal_gen/test/unit/realtime/test_lingbot_causal_denoising.py
+423 -5
View File
@@ -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):
python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py
+101 -1
View File
@@ -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)