From b0225a69dc5c7fbdd307298f314f9a9db3a00365 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 8 May 2026 01:18:42 +0800 Subject: [PATCH] [diffusion] optimize: precompute LTX2 guidance perturbation states (#24494) --- .../runtime/models/dits/ltx_2.py | 140 ++++++++++-------- 1 file changed, 81 insertions(+), 59 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py index 80c955611..96f513a92 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -64,32 +64,38 @@ def _ltx2_is_perturbed( return bool(value) -def _ltx2_batched_perturbation_mask( - perturbation_configs: tuple[dict[str, object], ...] | None, +def _ltx2_build_batched_perturbation_states( + perturbation_configs: tuple[dict[str, object], ...], key: str, - block_idx: int, + block_indices: tuple[int, ...], values: torch.Tensor, -) -> tuple[torch.Tensor | None, bool]: - if not perturbation_configs: - return None, False +) -> dict[int, tuple[torch.Tensor | None, bool]]: + mask_cache: dict[tuple[int, ...], torch.Tensor] = {} + states: dict[int, tuple[torch.Tensor | None, bool]] = {} + for block_idx in block_indices: + keep_values = [] + any_perturbed = False + all_perturbed = True + for config in perturbation_configs: + perturbed = _ltx2_is_perturbed(config, key, block_idx) + any_perturbed = any_perturbed or perturbed + all_perturbed = all_perturbed and perturbed + keep_values.append(0 if perturbed else 1) - mask = torch.ones( - (len(perturbation_configs),), device=values.device, dtype=values.dtype - ) - any_perturbed = False - all_perturbed = True - for batch_idx, config in enumerate(perturbation_configs): - perturbed = _ltx2_is_perturbed(config, key, block_idx) - any_perturbed = any_perturbed or perturbed - all_perturbed = all_perturbed and perturbed - if perturbed: - mask[batch_idx] = 0 - - if not any_perturbed: - return None, False - if all_perturbed: - return None, True - return mask.view(mask.numel(), *([1] * (values.ndim - 1))), False + if not any_perturbed: + states[block_idx] = (None, False) + elif all_perturbed: + states[block_idx] = (None, True) + else: + cache_key = tuple(keep_values) + mask = mask_cache.get(cache_key) + if mask is None: + mask = torch.tensor( + keep_values, device=values.device, dtype=values.dtype + ).view(len(keep_values), *([1] * (values.ndim - 1))) + mask_cache[cache_key] = mask + states[block_idx] = (mask, False) + return states @functools.lru_cache(maxsize=5) @@ -1811,6 +1817,46 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): # 5. Run blocks skip_video_self_attn_blocks = set(skip_video_self_attn_blocks or ()) skip_audio_self_attn_blocks = set(skip_audio_self_attn_blocks or ()) + video_self_attn_perturbation_states = None + audio_self_attn_perturbation_states = None + a2v_cross_attn_perturbation_states = None + v2a_cross_attn_perturbation_states = None + if perturbation_configs is not None: + block_indices = tuple( + getattr(block, "idx", -1) for block in self.transformer_blocks + ) + video_self_attn_perturbation_states = ( + _ltx2_build_batched_perturbation_states( + perturbation_configs, + "skip_video_self_attn_blocks", + block_indices, + hidden_states, + ) + ) + audio_self_attn_perturbation_states = ( + _ltx2_build_batched_perturbation_states( + perturbation_configs, + "skip_audio_self_attn_blocks", + block_indices, + audio_hidden_states, + ) + ) + a2v_cross_attn_perturbation_states = ( + _ltx2_build_batched_perturbation_states( + perturbation_configs, + "skip_a2v_cross_attn", + block_indices, + hidden_states, + ) + ) + v2a_cross_attn_perturbation_states = ( + _ltx2_build_batched_perturbation_states( + perturbation_configs, + "skip_v2a_cross_attn", + block_indices, + audio_hidden_states, + ) + ) for block in self.transformer_blocks: block_idx = getattr(block, "idx", -1) video_self_attn_perturbation_mask = None @@ -1823,45 +1869,21 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): skip_v2a_cross_attn = disable_v2a_cross_attn if perturbation_configs is not None: if not skip_video_self_attn: - ( - video_self_attn_perturbation_mask, - skip_video_self_attn, - ) = _ltx2_batched_perturbation_mask( - perturbation_configs, - "skip_video_self_attn_blocks", - block_idx, - hidden_states, - ) + assert video_self_attn_perturbation_states is not None + state = video_self_attn_perturbation_states[block_idx] + video_self_attn_perturbation_mask, skip_video_self_attn = state if not skip_audio_self_attn: - ( - audio_self_attn_perturbation_mask, - skip_audio_self_attn, - ) = _ltx2_batched_perturbation_mask( - perturbation_configs, - "skip_audio_self_attn_blocks", - block_idx, - audio_hidden_states, - ) + assert audio_self_attn_perturbation_states is not None + state = audio_self_attn_perturbation_states[block_idx] + audio_self_attn_perturbation_mask, skip_audio_self_attn = state if not skip_a2v_cross_attn: - ( - a2v_cross_attn_perturbation_mask, - skip_a2v_cross_attn, - ) = _ltx2_batched_perturbation_mask( - perturbation_configs, - "skip_a2v_cross_attn", - block_idx, - hidden_states, - ) + assert a2v_cross_attn_perturbation_states is not None + state = a2v_cross_attn_perturbation_states[block_idx] + a2v_cross_attn_perturbation_mask, skip_a2v_cross_attn = state if not skip_v2a_cross_attn: - ( - v2a_cross_attn_perturbation_mask, - skip_v2a_cross_attn, - ) = _ltx2_batched_perturbation_mask( - perturbation_configs, - "skip_v2a_cross_attn", - block_idx, - audio_hidden_states, - ) + assert v2a_cross_attn_perturbation_states is not None + state = v2a_cross_attn_perturbation_states[block_idx] + v2a_cross_attn_perturbation_mask, skip_v2a_cross_attn = state hidden_states, audio_hidden_states = block( hidden_states, audio_hidden_states,