minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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)
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,