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 5df820e65..11c17d9e0 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py @@ -80,6 +80,7 @@ from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) +from sglang.multimodal_gen.runtime.realtime.causal_state import RealtimeCausalDiTState from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.srt.utils import add_prefix @@ -180,6 +181,7 @@ class LingBotWorldCausalSelfAttention(CausalWanSelfAttention): kv_cache: CausalSelfAttentionKVCache | None = None, current_start: int = 0, cache_start: int | None = None, + update_cache_only: bool = False, ): cos, sin = freqs_cis[:2] cos_sin_cache = freqs_cis[2] if len(freqs_cis) > 2 else None @@ -243,6 +245,8 @@ class LingBotWorldCausalSelfAttention(CausalWanSelfAttention): current_chunk_start=current_start, debug_name="LingBot KV cache", ) + if update_cache_only: + return v attn_impl = self.ulysses_attn if sequence_shard_enabled else self.attn x = attn_impl( roped_query, @@ -926,6 +930,7 @@ class CausalLingBotWorldTransformerBlock(CausalWanTransformerBlock): current_start: int = 0, cache_start: int | None = None, c2ws_plucker_emb: torch.Tensor | None = None, + update_cache_only: bool = False, ) -> torch.Tensor: if hidden_states.dim() == 4: hidden_states = hidden_states.squeeze(1) @@ -963,7 +968,10 @@ class CausalLingBotWorldTransformerBlock(CausalWanTransformerBlock): kv_cache, current_start, cache_start, + update_cache_only=update_cache_only, ) + if update_cache_only: + return hidden_states attn_output = attn_output.flatten(2) attn_output, _ = self.to_out(attn_output) attn_output = attn_output.squeeze(1) @@ -1113,6 +1121,10 @@ class CausalLingBotWorldTransformer3DModel(CausalWanTransformer3DModel): def _get_request_cache(forward_batch, name: str) -> dict | None: if forward_batch is None: return None + session = getattr(forward_batch, "session", None) + if session is not None: + state = session.get_or_create_state(RealtimeCausalDiTState) + return state.runtime_cache.setdefault(name, {}) extra = getattr(forward_batch, "extra", None) if extra is None: return None @@ -1179,6 +1191,49 @@ class CausalLingBotWorldTransformer3DModel(CausalWanTransformer3DModel): cache[cache_key] = freqs_cis return freqs_cis + def _prepare_cached_rope_for_sequence_shard( + self, + *, + forward_batch, + local_seq_len: int, + token_start: int, + frame_stride: int, + post_patch_width: int, + device: torch.device, + ) -> tuple[torch.Tensor, ...]: + cache = self._get_request_cache(forward_batch, "lingbot_sequence_shard_rope") + cache_key = ( + local_seq_len, + token_start, + frame_stride, + post_patch_width, + device.type, + device.index, + ) + if cache is not None and cache_key in cache: + return cache[cache_key] + + freqs_cos, freqs_sin = self._compute_rope_for_sequence_shard_with_offset( + local_seq_len, + token_start, + frame_stride, + post_patch_width, + device, + ) + freqs_cos = freqs_cos.float() + freqs_sin = freqs_sin.float() + freqs_cis: tuple[torch.Tensor, ...] = (freqs_cos, freqs_sin) + if _is_cuda: + freqs_cis = ( + freqs_cos, + freqs_sin, + torch.cat([freqs_cos.contiguous(), freqs_sin.contiguous()], dim=-1), + ) + if cache is not None: + cache.clear() + cache[cache_key] = freqs_cis + return freqs_cis + def _prepare_condition_embeddings( self, *, @@ -1187,14 +1242,48 @@ class CausalLingBotWorldTransformer3DModel(CausalWanTransformer3DModel): encoder_hidden_states_image: torch.Tensor | None, crossattn_cache: list[CrossAttentionKVCache] | None, ): + forward_batch = get_forward_context().forward_batch + temb, timestep_proj = self._prepare_cached_time_embeddings( + timestep=timestep, + forward_batch=forward_batch, + ) if self._all_crossattn_caches_initialized(crossattn_cache): - temb = self.condition_embedder.time_embedder(timestep.flatten()) - timestep_proj = self.condition_embedder.time_modulation(temb) return temb, timestep_proj, encoder_hidden_states, None - return self.condition_embedder( - timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image + encoder_hidden_states = self.condition_embedder.text_embedder( + encoder_hidden_states ) + if encoder_hidden_states_image is not None: + assert self.condition_embedder.image_embedder is not None + encoder_hidden_states_image = self.condition_embedder.image_embedder( + encoder_hidden_states_image + ) + + return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image + + def _prepare_cached_time_embeddings( + self, + *, + timestep: torch.LongTensor, + forward_batch, + ) -> tuple[torch.Tensor, torch.Tensor]: + cache = self._get_request_cache(forward_batch, "lingbot_time_embeddings") + current_timestep = get_forward_context().current_timestep + cache_key = ( + current_timestep, + tuple(timestep.shape), + timestep.dtype, + timestep.device.type, + timestep.device.index, + ) + if cache is not None and cache_key in cache: + return cache[cache_key] + + temb = self.condition_embedder.time_embedder(timestep.flatten()) + timestep_proj = self.condition_embedder.time_modulation(temb) + if cache is not None: + cache[cache_key] = (temb, timestep_proj) + return temb, timestep_proj def forward( self, @@ -1274,22 +1363,14 @@ class CausalLingBotWorldTransformer3DModel(CausalWanTransformer3DModel): ) frame_stride = post_patch_height * post_patch_width token_start = start_frame * frame_stride + sum(seq_shard_splits[:sp_rank]) - freqs_cos, freqs_sin = self._compute_rope_for_sequence_shard_with_offset( - local_seq_len, - token_start, - frame_stride, - post_patch_width, - hidden_states.device, + freqs_cis = self._prepare_cached_rope_for_sequence_shard( + forward_batch=forward_batch, + local_seq_len=local_seq_len, + token_start=token_start, + frame_stride=frame_stride, + post_patch_width=post_patch_width, + device=hidden_states.device, ) - freqs_cos = freqs_cos.float() - freqs_sin = freqs_sin.float() - freqs_cis = (freqs_cos, freqs_sin) - if _is_cuda: - freqs_cis = ( - freqs_cos, - freqs_sin, - torch.cat([freqs_cos.contiguous(), freqs_sin.contiguous()], dim=-1), - ) temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = ( self._prepare_condition_embeddings( @@ -1324,22 +1405,23 @@ class CausalLingBotWorldTransformer3DModel(CausalWanTransformer3DModel): current_start=current_start, cache_start=cache_start, c2ws_plucker_emb=c2ws_plucker_emb, + update_cache_only=skip_final_projection + and block_index == len(self.blocks) - 1, ) if skip_final_projection: return hidden_states + temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2) + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) + hidden_states = self.norm_out(hidden_states, shift, scale) + hidden_states = self.proj_out(hidden_states) if sequence_shard_enabled: hidden_states = _sequence_all_gather_varlen( hidden_states.contiguous(), list(forward_batch.sequence_shard_splits), get_sp_group().device_group, ) - - temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2) - shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) - hidden_states = self.norm_out(hidden_states, shift, scale) - hidden_states = self.proj_out(hidden_states) hidden_states = hidden_states.reshape( batch_size, post_patch_num_frames, 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 3fa583d22..64b5fab41 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 @@ -190,7 +190,7 @@ class LingBotWorldCausalDMDDenoisingStage(CausalDMDDenoisingStage): enabled=autocast_enabled, ), set_forward_context( - current_timestep=0, + current_timestep=-1, attn_metadata=attn_metadata, forward_batch=batch, ), diff --git a/python/sglang/multimodal_gen/runtime/realtime/causal_state.py b/python/sglang/multimodal_gen/runtime/realtime/causal_state.py index 89ebb98fd..d7de5eb16 100644 --- a/python/sglang/multimodal_gen/runtime/realtime/causal_state.py +++ b/python/sglang/multimodal_gen/runtime/realtime/causal_state.py @@ -10,11 +10,13 @@ class RealtimeCausalDiTState(BaseRealtimeState): super().__init__() self.kv_cache = None self.crossattn_cache = None + self.runtime_cache: dict = {} self.current_chunk_start_frame: int = 0 self.chunk_idx: int = 0 def dispose(self) -> None: self.kv_cache = None self.crossattn_cache = None + self.runtime_cache.clear() self.current_chunk_start_frame = 0 self.chunk_idx = 0 diff --git a/scripts/ci/utils/diffusion/diffusion_case_parser.py b/scripts/ci/utils/diffusion/diffusion_case_parser.py index ef7ce822b..74159e668 100755 --- a/scripts/ci/utils/diffusion/diffusion_case_parser.py +++ b/scripts/ci/utils/diffusion/diffusion_case_parser.py @@ -80,6 +80,21 @@ class DiffusionTestCaseVisitor(ast.NodeVisitor): def __init__(self): self.cases: Dict[str, List[str]] = {} # list_name -> [case_id, ...] + self.factory_case_ids: Dict[str, str] = {} + + def visit_Module(self, node: ast.Module): + for stmt in node.body: + if not isinstance(stmt, ast.FunctionDef): + continue + case_id = self._extract_factory_case_id(stmt) + if case_id: + self.factory_case_ids[stmt.name] = case_id + + for stmt in node.body: + if isinstance(stmt, ast.Expr): + self._process_expr(stmt.value) + + self.generic_visit(node) def visit_Assign(self, node: ast.Assign): self._process_assignment(node.targets, node.value) @@ -121,6 +136,26 @@ class DiffusionTestCaseVisitor(ast.NodeVisitor): lhs_case_ids = self.cases.get(target.id, []) self.cases[target.id] = [*lhs_case_ids, *rhs_case_ids] + def _process_expr(self, node: ast.AST): + """Process list mutation calls such as `ONE_GPU_CASES.append(...)`.""" + if not isinstance(node, ast.Call): + return + if not isinstance(node.func, ast.Attribute): + return + if node.func.attr != "append": + return + if not isinstance(node.func.value, ast.Name): + return + list_name = node.func.value.id + if list_name not in CASE_LIST_TO_SUITE: + return + if len(node.args) != 1: + return + + case_id = self._extract_case_id_from_call(node.args[0]) + if case_id: + self.cases.setdefault(list_name, []).append(case_id) + def _extract_case_ids(self, node: ast.AST) -> Optional[List[str]]: """Extract case IDs from a supported expression.""" if isinstance(node, ast.List): @@ -167,9 +202,20 @@ class DiffusionTestCaseVisitor(ast.NodeVisitor): }: if node.args and isinstance(node.args[0], ast.Constant): return node.args[0].value + if isinstance(node.func, ast.Name) and not node.args: + return self.factory_case_ids.get(node.func.id) return None + def _extract_factory_case_id(self, node: ast.FunctionDef) -> Optional[str]: + for child in ast.walk(node): + if not isinstance(child, ast.Return) or child.value is None: + continue + case_id = self._extract_case_id_from_call(child.value) + if case_id: + return case_id + return None + def resolve_case_config_path(repo_root: Path, run_suite_path: Path) -> Path: """