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

[diffusion] optimize: precompute LTX2 guidance perturbation states (#24494)

Browse Source
This commit is contained in:
Mick
2026-05-08 01:18:42 +08:00
committed by GitHub
parent c18879054f
commit b0225a69dc
1 changed files with 81 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py
+81 -59
View File
@@ -64,32 +64,38 @@ def _ltx2_is_perturbed(
return bool(value) return bool(value)
def _ltx2_batched_perturbation_mask( def _ltx2_build_batched_perturbation_states(
perturbation_configs: tuple[dict[str, object], ...] | None, perturbation_configs: tuple[dict[str, object], ...],
key: str, key: str,
block_idx: int, block_indices: tuple[int, ...],
values: torch.Tensor, values: torch.Tensor,
) -> tuple[torch.Tensor | None, bool]: ) -> dict[int, tuple[torch.Tensor | None, bool]]:
if not perturbation_configs: mask_cache: dict[tuple[int, ...], torch.Tensor] = {}
return None, False 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( if not any_perturbed:
(len(perturbation_configs),), device=values.device, dtype=values.dtype states[block_idx] = (None, False)
) elif all_perturbed:
any_perturbed = False states[block_idx] = (None, True)
all_perturbed = True else:
for batch_idx, config in enumerate(perturbation_configs): cache_key = tuple(keep_values)
perturbed = _ltx2_is_perturbed(config, key, block_idx) mask = mask_cache.get(cache_key)
any_perturbed = any_perturbed or perturbed if mask is None:
all_perturbed = all_perturbed and perturbed mask = torch.tensor(
if perturbed: keep_values, device=values.device, dtype=values.dtype
mask[batch_idx] = 0 ).view(len(keep_values), *([1] * (values.ndim - 1)))
mask_cache[cache_key] = mask
if not any_perturbed: states[block_idx] = (mask, False)
return None, False return states
if all_perturbed:
return None, True
return mask.view(mask.numel(), *([1] * (values.ndim - 1))), False
@functools.lru_cache(maxsize=5) @functools.lru_cache(maxsize=5)
@@ -1811,6 +1817,46 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
# 5. Run blocks # 5. Run blocks
skip_video_self_attn_blocks = set(skip_video_self_attn_blocks or ()) skip_video_self_attn_blocks = set(skip_video_self_attn_blocks or ())
skip_audio_self_attn_blocks = set(skip_audio_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: for block in self.transformer_blocks:
block_idx = getattr(block, "idx", -1) block_idx = getattr(block, "idx", -1)
video_self_attn_perturbation_mask = None video_self_attn_perturbation_mask = None
@@ -1823,45 +1869,21 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
skip_v2a_cross_attn = disable_v2a_cross_attn skip_v2a_cross_attn = disable_v2a_cross_attn
if perturbation_configs is not None: if perturbation_configs is not None:
if not skip_video_self_attn: if not skip_video_self_attn:
( assert video_self_attn_perturbation_states is not None
video_self_attn_perturbation_mask, state = video_self_attn_perturbation_states[block_idx]
skip_video_self_attn, video_self_attn_perturbation_mask, skip_video_self_attn = state
) = _ltx2_batched_perturbation_mask(
perturbation_configs,
"skip_video_self_attn_blocks",
block_idx,
hidden_states,
)
if not skip_audio_self_attn: if not skip_audio_self_attn:
( assert audio_self_attn_perturbation_states is not None
audio_self_attn_perturbation_mask, state = audio_self_attn_perturbation_states[block_idx]
skip_audio_self_attn, audio_self_attn_perturbation_mask, skip_audio_self_attn = state
) = _ltx2_batched_perturbation_mask(
perturbation_configs,
"skip_audio_self_attn_blocks",
block_idx,
audio_hidden_states,
)
if not skip_a2v_cross_attn: if not skip_a2v_cross_attn:
( assert a2v_cross_attn_perturbation_states is not None
a2v_cross_attn_perturbation_mask, state = a2v_cross_attn_perturbation_states[block_idx]
skip_a2v_cross_attn, a2v_cross_attn_perturbation_mask, skip_a2v_cross_attn = state
) = _ltx2_batched_perturbation_mask(
perturbation_configs,
"skip_a2v_cross_attn",
block_idx,
hidden_states,
)
if not skip_v2a_cross_attn: if not skip_v2a_cross_attn:
( assert v2a_cross_attn_perturbation_states is not None
v2a_cross_attn_perturbation_mask, state = v2a_cross_attn_perturbation_states[block_idx]
skip_v2a_cross_attn, v2a_cross_attn_perturbation_mask, skip_v2a_cross_attn = state
) = _ltx2_batched_perturbation_mask(
perturbation_configs,
"skip_v2a_cross_attn",
block_idx,
audio_hidden_states,
)
hidden_states, audio_hidden_states = block( hidden_states, audio_hidden_states = block(
hidden_states, hidden_states,
audio_hidden_states, audio_hidden_states,