diff --git a/tests/diffusion/models/ltx2/test_ltx2_diffvae_attention.py b/tests/diffusion/models/ltx2/test_ltx2_diffvae_attention.py new file mode 100644 index 000000000..6cdaf197d --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_diffvae_attention.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Native FNA against an independent, test-only explicit neighborhood mask.""" + +import pytest +import torch +from torch.nn.attention import SDPBackend, sdpa_kernel + +from vllm_omni.diffusion.models.ltx2.ops import resolve_ltx2_vae_operators +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion] + + +def _selected_fna(): + if not current_omni_platform.is_available(): + pytest.skip("No accelerator is available") + operators = resolve_ltx2_vae_operators(current_omni_platform.get_torch_device()) + if operators is None or operators.fna is None: + pytest.skip("LTX DiffVAE FNA is unavailable for this device or installation") + return operators.fna + + +def _load_fna_ops(): + return pytest.importorskip("vllm_omni.diffusion.models.ltx2.ops.fna", exc_type=ImportError) + + +def _pack(x, tile): + b, t, h, w, heads, dim = x.shape + tt, th, tw = tile + nt, nh, nw = (t + tt - 1) // tt, (h + th - 1) // th, (w + tw - 1) // tw + padded = x.new_zeros((b, nt * tt, nh * th, nw * tw, heads, dim)) + padded[:, :t, :h, :w] = x + # Outer and inner spatial tiles both flatten in W,H,T order. + return ( + padded.view(b, nt, tt, nh, th, nw, tw, heads, dim) + .permute(0, 5, 3, 1, 6, 4, 2, 7, 8) + .reshape(b, nt * nh * nw * tt * th * tw, heads, dim) + ) + + +def _unpack(x, shape): + t, h, w = shape + b, _, heads, dim = x.shape + nt, nh, nw = ((n + 3) // 4 for n in shape) + return ( + x.view(b, nw, nh, nt, 4, 4, 4, heads, dim) + .permute(0, 3, 6, 2, 5, 1, 4, 7, 8) + .reshape(b, nt * 4, nh * 4, nw * 4, heads, dim)[:, :t, :h, :w] + ) + + +@pytest.mark.parametrize("batch,shape", [(1, (11, 11, 11)), (2, (13, 12, 17))]) +@torch.inference_mode() +def test_fna_matches_explicit_inward_shift_mask(batch, shape): + fna = _selected_fna() + generator = torch.Generator(device="cuda").manual_seed(42) + q, k, v = [ + torch.randn(batch, *shape, 4, 64, device="cuda", dtype=torch.bfloat16, generator=generator) for _ in range(3) + ] + q = q * 0.125 + coordinates = torch.stack( + torch.meshgrid(*(torch.arange(n, device="cuda") for n in shape), indexing="ij"), dim=-1 + ).reshape(-1, 3) + start = (coordinates - 5).clamp_min(0).minimum(torch.tensor(shape, device="cuda") - 11) + mask = ((coordinates[None] >= start[:, None]) & (coordinates[None] < start[:, None] + 11)).all(-1) + with sdpa_kernel(SDPBackend.MATH): + expected = ( + torch.nn.functional.scaled_dot_product_attention( + *(x.flatten(1, 3).transpose(1, 2).float() for x in (q, k, v)), attn_mask=mask, scale=1.0 + ) + .transpose(1, 2) + .reshape(q.shape) + ) + actual = _unpack( + fna(_pack(q, (4, 4, 4)), _pack(k, (4, 4, 8)), _pack(v, (4, 4, 8)), shape=shape, window=(11, 11, 11)), + shape, + ).float() + assert actual.isfinite().all() + assert float((actual - expected).norm() / expected.norm()) < 0.005 + torch.testing.assert_close(actual, expected, atol=0.008, rtol=0.025) + + +def test_fna_rejects_cpu_without_fallback(): + fna_ops = _load_fna_ops() + x = torch.zeros(1, 64, 4, 64, dtype=torch.bfloat16) + with pytest.raises(ValueError, match="CUDA BF16"): + fna_ops.fna3d_tilelang(x, x, x, shape=(11, 11, 11), window=(11, 11, 11)) diff --git a/tests/diffusion/models/ltx2/test_ltx2_diffvae_ops.py b/tests/diffusion/models/ltx2/test_ltx2_diffvae_ops.py new file mode 100644 index 000000000..a66a7f7c8 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_diffvae_ops.py @@ -0,0 +1,409 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import subprocess +import sys + +import pytest +import torch + +from vllm_omni.diffusion.models.ltx2.ops import qk_rms_norm as qk_rms_norm_ops +from vllm_omni.diffusion.models.ltx2.ops import residual_adaln as residual_adaln_ops +from vllm_omni.diffusion.models.ltx2.ops import resolve_ltx2_vae_operators +from vllm_omni.diffusion.models.ltx2.ops import swiglu as swiglu_ops +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion] + +_DIM_SPLIT = (16, 24, 24) + + +def _selected_operators(): + if not current_omni_platform.is_available(): + pytest.skip("No accelerator is available") + operators = resolve_ltx2_vae_operators(current_omni_platform.get_torch_device()) + if operators is None: + pytest.skip("No LTX DiffVAE operators are available for this device") + return operators + + +def _selected_fna(): + operators = _selected_operators() + if operators.fna is None: + pytest.skip("LTX DiffVAE FNA is unavailable for this device or installation") + return operators.fna + + +def _load_fna_ops(): + return pytest.importorskip("vllm_omni.diffusion.models.ltx2.ops.fna", exc_type=ImportError) + + +def test_diffvae_model_and_tests_do_not_import_tilelang_backend() -> None: + # Use a fresh interpreter: reloading the model in this process would leave + # already-imported distributed subclasses pointing at the old base class. + result = subprocess.run( + [ + sys.executable, + "-c", + "import sys; " + "sys.modules['vllm_omni.diffusion.models.ltx2.ops.fna'] = None; " + "sys.modules['vllm.tilelang_utils'] = None; " + "from vllm_omni.diffusion.models.ltx2.vae.decoder import LTX2VideoDiffusionDecoderModel; " + "import tests.diffusion.models.ltx2.test_ltx2_diffvae_attention; " + "import tests.diffusion.models.ltx2.test_ltx2_diffvae_ops", + ], + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +def test_diffvae_fna_tile_map_covers_exact_neighborhood() -> None: + fna_ops = _load_fna_ops() + problem = fna_ops.Problem(13, 14, 17, (11, 11, 11)) + indices, counts = fna_ops.build_kv_tile_map(problem) + q_grid_t, q_grid_h, _ = problem.q_grid + kv_grid_t, kv_grid_h, _ = problem.kv_grid + + for q_block in range(indices.shape[0]): + selected = set(indices[q_block, : counts[q_block]].tolist()) + q_outer_t = q_block % q_grid_t + q_outer_h = (q_block // q_grid_t) % q_grid_h + q_outer_w = q_block // (q_grid_t * q_grid_h) + for q_inner in range(fna_ops.Q_TOKENS): + query = ( + q_outer_t * 4 + q_inner % 4, + q_outer_h * 4 + (q_inner // 4) % 4, + q_outer_w * 4 + q_inner // 16, + ) + if any(index >= length for index, length in zip(query, problem.shape, strict=True)): + continue + starts = tuple( + fna_ops._window_start(index, window, length) + for index, window, length in zip(query, problem.window, problem.shape, strict=True) + ) + for key_t in range(starts[0], starts[0] + problem.window[0]): + for key_h in range(starts[1], starts[1] + problem.window[1]): + for key_w in range(starts[2], starts[2] + problem.window[2]): + key_block = ((key_w // 8) * kv_grid_h + key_h // 4) * kv_grid_t + key_t // 4 + assert key_block in selected + + +def test_diffvae_fna_metadata_uses_static_kernel_shapes() -> None: + fna_ops = _load_fna_ops() + for shape in ((11, 11, 11), (13, 14, 17), (66, 96, 128), (79, 192, 192)): + problem = fna_ops.Problem(*shape, (11, 11, 11)) + _, t_patterns, h_patterns, w_patterns = fna_ops.build_kv_tile_metadata(problem) + assert {table.numel() for table in (t_patterns, h_patterns, w_patterns)} == {fna_ops.PATTERN_TABLE_SIZE} + assert {limit for _, _, limit in fna_ops.build_kv_tile_buckets(problem)} <= set(fna_ops.KV_BUCKET_LIMITS) + for ids, metadata, limit in fna_ops.build_kv_tile_buckets(problem): + assert metadata.shape == (ids.numel(), limit) + + +def test_diffvae_fna_rejects_oversized_packed_tile_grid() -> None: + fna_ops = _load_fna_ops() + with pytest.raises(ValueError, match="15-bit metadata"): + fna_ops.Problem(128, 256, 256, (11, 11, 11)) + + +def test_diffvae_fna_propagates_launch_failures( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fna_ops = _load_fna_ops() + fna = _selected_fna() + problem = fna_ops.Problem(11, 11, 11, (11, 11, 11)) + query = torch.randn(1, problem.q_length, fna_ops.HEADS, 64, device="cuda", dtype=torch.bfloat16) + key = torch.randn(1, problem.kv_length, fna_ops.HEADS, 64, device="cuda", dtype=torch.bfloat16) + value = torch.randn_like(key) + calls = 0 + + def fail_launch(*args, **kwargs): + nonlocal calls + calls += 1 + raise RuntimeError("injected launch failure") + + monkeypatch.setattr(fna_ops, "_fna3d_tilelang", fail_launch) + with torch.inference_mode(): + for _ in range(2): + with pytest.raises(RuntimeError, match="injected launch failure"): + fna( + query, + key, + value, + shape=problem.shape, + window=problem.window, + ) + assert calls == 2 + + +@pytest.fixture +def isolated_runtime_state(): + targets = ( + (qk_rms_norm_ops, "_FAILED_ROPE_KEYS"), + (qk_rms_norm_ops, "_VERIFIED_ROPE_KEYS"), + (swiglu_ops, "_FAILED_KEYS"), + (swiglu_ops, "_VERIFIED_KEYS"), + (residual_adaln_ops, "_FAILED_ADALN_KEYS"), + (residual_adaln_ops, "_VERIFIED_ADALN_KEYS"), + (residual_adaln_ops, "_FAILED_ADD_DEVICES"), + (residual_adaln_ops, "_VERIFIED_ADD_DEVICES"), + ) + saved = [set(getattr(module, name)) for module, name in targets] + for module, name in targets: + getattr(module, name).clear() + try: + yield + finally: + for (module, name), state in zip(targets, saved, strict=True): + target = getattr(module, name) + target.clear() + target.update(state) + + +@pytest.mark.parametrize( + ("dim", "rows"), + [ + (256, 257), + (256, 16385), + (256, 131075), + (512, 65539), + (1024, 32771), + (2048, 16387), + ], +) +def test_diffvae_swiglu_tiled_is_bit_exact(dim: int, rows: int, isolated_runtime_state) -> None: + from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoVaeSwiGLU + + operators = _selected_operators() + generator = torch.Generator(device="cuda").manual_seed(20260829) + hidden_states = torch.randn((1, 1, 1, rows, dim), device="cuda", dtype=torch.bfloat16, generator=generator) + gate_weight = torch.randn((4 * dim, dim), device="cuda", dtype=torch.bfloat16, generator=generator) + up_weight = torch.randn_like(gate_weight) + down_weight = torch.randn((dim, 4 * dim), device="cuda", dtype=torch.bfloat16, generator=generator) + with torch.device("meta"): + reference = LTX2VideoVaeSwiGLU(dim, 4 * dim) + reference.w_gate.weight = torch.nn.Parameter(gate_weight, requires_grad=False) + reference.w_up.weight = torch.nn.Parameter(up_weight, requires_grad=False) + reference.w_down.weight = torch.nn.Parameter(down_weight, requires_grad=False) + + with torch.inference_mode(): + expected = reference(hidden_states) + actual = operators.swiglu(hidden_states, gate_weight, up_weight, down_weight) + + assert actual is not None + assert torch.equal(actual, expected) + + +@pytest.mark.parametrize( + ("dim", "expected_rows"), + [(256, 131072), (512, 65536), (1024, 32768), (2048, 16384)], +) +def test_diffvae_swiglu_tile_rows_keep_a_constant_hidden_budget(dim: int, expected_rows: int) -> None: + assert swiglu_ops._tile_rows(1_000_000, 4 * dim) == expected_rows + + +def test_diffvae_swiglu_tiled_permanently_falls_back_after_failure( + monkeypatch: pytest.MonkeyPatch, + isolated_runtime_state, +) -> None: + operators = _selected_operators() + hidden_states = torch.randn(1, 1, 1, 17, 256, device="cuda", dtype=torch.bfloat16) + gate_weight = torch.randn(1024, 256, device="cuda", dtype=torch.bfloat16) + up_weight = torch.randn_like(gate_weight) + down_weight = torch.randn(256, 1024, device="cuda", dtype=torch.bfloat16) + calls = 0 + + def fail_launch(*args, **kwargs): + nonlocal calls + calls += 1 + raise RuntimeError("injected launch failure") + + monkeypatch.setattr(swiglu_ops, "_launch", fail_launch) + with torch.inference_mode(): + assert operators.swiglu(hidden_states, gate_weight, up_weight, down_weight) is None + assert operators.swiglu(hidden_states, gate_weight, up_weight, down_weight) is None + + assert calls == 1 + + +@pytest.mark.parametrize("residual_count", [1, 2]) +def test_diffvae_residual_rms_norm_modulate_is_bit_exact( + residual_count: int, + isolated_runtime_state, +) -> None: + operators = _selected_operators() + generator = torch.Generator(device="cuda").manual_seed(20260829) + shape = (2, 3, 7, 7, 256) + x = torch.randn(shape, device="cuda", dtype=torch.bfloat16, generator=generator) + residual_a = torch.randn_like(x) + residual_b = torch.randn_like(x) if residual_count == 2 else None + norm_weight = torch.randn((256,), device="cuda", dtype=torch.bfloat16, generator=generator) + scale = torch.randn((2, 1, 1, 1, 256), device="cuda", dtype=torch.bfloat16, generator=generator) + shift = torch.randn_like(scale) + hidden_states = x + residual_a + if residual_b is not None: + hidden_states = hidden_states + residual_b + expected = torch.nn.functional.rms_norm(hidden_states, (256,), norm_weight, eps=1e-6) * (1 + scale) + shift + + with torch.inference_mode(): + actual = operators.residual_norm( + x, + residual_a, + residual_b, + norm_weight, + scale, + shift, + 1e-6, + ) + + assert actual is not None + assert torch.equal(actual, expected) + + +def test_diffvae_residual_add3_is_bit_exact(isolated_runtime_state) -> None: + operators = _selected_operators() + generator = torch.Generator(device="cuda").manual_seed(20260829) + tensors = [ + torch.randn((2, 3, 7, 7, 256), device="cuda", dtype=torch.bfloat16, generator=generator) for _ in range(4) + ] + expected = tensors[0] + tensors[1] + expected = expected + tensors[2] + expected = expected + tensors[3] + + with torch.inference_mode(): + actual = operators.residual_add(*tensors) + + assert actual is not None + assert torch.equal(actual, expected) + + +def test_diffvae_residual_adaln_permanently_falls_back_after_failure( + monkeypatch: pytest.MonkeyPatch, + isolated_runtime_state, +) -> None: + operators = _selected_operators() + shape = (1, 3, 7, 7, 256) + x = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + residual = torch.randn_like(x) + weight = torch.randn(256, device="cuda", dtype=torch.bfloat16) + scale = torch.randn(1, 1, 1, 1, 256, device="cuda", dtype=torch.bfloat16) + shift = torch.randn_like(scale) + calls = 0 + + def fail_launch(*args, **kwargs): + nonlocal calls + calls += 1 + raise RuntimeError("injected launch failure") + + monkeypatch.setattr(residual_adaln_ops, "_launch_adaln", fail_launch) + with torch.inference_mode(): + for _ in range(2): + assert operators.residual_norm(x, residual, None, weight, scale, shift, 1e-6) is None + + assert calls == 1 + + +def _tables( + frames: int, + height: int, + width: int, + dim_split: tuple[int, int, int] = _DIM_SPLIT, + base: float = 10000.0, +) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: + tables = [] + for length, dim in zip((frames, height, width), dim_split, strict=True): + exponents = torch.arange(0, dim, 2, dtype=torch.float64, device="cuda") / dim + inv_freqs = (1.0 / base**exponents).to(torch.float32) + positions = torch.arange(length, dtype=torch.float32, device="cuda") + angles = positions[:, None] * inv_freqs[None, :] + tables.append((angles.cos(), angles.sin())) + return tuple(tables) + + +def _apply_rope_reference( + hidden_states: torch.Tensor, + tables: tuple[tuple[torch.Tensor, torch.Tensor], ...], + dim_split: tuple[int, int, int] = _DIM_SPLIT, +) -> torch.Tensor: + outputs = [] + offset = 0 + for axis, (dim, (cos, sin)) in enumerate(zip(dim_split, tables, strict=True), 1): + chunk = hidden_states[..., offset : offset + dim] + pairs = chunk.reshape(*chunk.shape[:-1], dim // 2, 2) + even = pairs[..., 0].float() + odd = pairs[..., 1].float() + shape = [1, 1, 1, 1, 1, dim // 2] + shape[axis] = cos.shape[0] + cos = cos.reshape(shape) + sin = sin.reshape(shape) + rotated = torch.stack([even * cos - odd * sin, even * sin + odd * cos], dim=-1) + outputs.append(rotated.reshape(chunk.shape).to(hidden_states.dtype)) + offset += dim + return torch.cat(outputs, dim=-1) + + +@pytest.mark.parametrize( + ("frames", "height", "width", "heads"), + [(3, 7, 7, 1), (31, 136, 192, 4)], +) +def test_diffvae_qk_rms_norm_scale_rope_3d_is_bit_exact( + frames: int, + height: int, + width: int, + heads: int, + isolated_runtime_state, +) -> None: + operators = _selected_operators() + generator = torch.Generator(device="cuda").manual_seed(20260829) + shape = (1, frames, height, width, heads, 64) + query = torch.randn(shape, device="cuda", dtype=torch.bfloat16, generator=generator) + key = torch.randn_like(query) + query_weight = torch.randn((64,), device="cuda", dtype=torch.bfloat16, generator=generator) + key_weight = torch.randn((64,), device="cuda", dtype=torch.bfloat16, generator=generator) + tables = _tables(frames, height, width) + expected = ( + _apply_rope_reference(torch.nn.functional.rms_norm(query, (64,), query_weight, 1e-6) * 0.125, tables), + _apply_rope_reference(torch.nn.functional.rms_norm(key, (64,), key_weight, 1e-6), tables), + ) + + with torch.inference_mode(): + actual = operators.qk_norm_rope( + query, + key, + query_weight, + key_weight, + 1e-6, + 0.125, + _DIM_SPLIT, + 10000.0, + ) + + assert actual is not None + assert torch.equal(actual[0], expected[0]) + assert torch.equal(actual[1], expected[1]) + + +def test_diffvae_qk_rms_norm_scale_rope_3d_permanently_falls_back_after_failure( + monkeypatch: pytest.MonkeyPatch, + isolated_runtime_state, +) -> None: + operators = _selected_operators() + query = torch.randn(1, 3, 7, 7, 1, 64, device="cuda", dtype=torch.bfloat16) + key = torch.randn_like(query) + weight = torch.ones(64, device="cuda", dtype=torch.bfloat16) + calls = 0 + + def fail_launch(*args, **kwargs): + nonlocal calls + calls += 1 + raise RuntimeError("injected launch failure") + + monkeypatch.setattr(qk_rms_norm_ops, "_launch_combined", fail_launch) + with torch.inference_mode(): + for _ in range(2): + assert operators.qk_norm_rope(query, key, weight, weight, 1e-6, 0.125, _DIM_SPLIT, 10000.0) is None + + assert calls == 1 diff --git a/tests/diffusion/models/ltx2/test_ltx2_ops_platform.py b/tests/diffusion/models/ltx2/test_ltx2_ops_platform.py new file mode 100644 index 000000000..6ffff3136 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_ops_platform.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import sys +from dataclasses import dataclass + +import pytest +import torch +from vllm.platforms.interface import DeviceCapability + +from vllm_omni.diffusion.models.ltx2 import ops as ltx2_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@dataclass(frozen=True) +class _TensorProbe: + is_cuda: bool + device: torch.device + + +@pytest.fixture(autouse=True) +def clear_platform_cache(): + ltx2_platform._cuda_compute_capability.cache_clear() + try: + yield + finally: + ltx2_platform._cuda_compute_capability.cache_clear() + + +@pytest.mark.parametrize( + ("capability", "triton", "tilelang", "fusion_expected", "fna_expected"), + [ + (80, True, True, False, False), + (90, True, True, True, True), + (90, True, False, True, False), + (90, False, True, False, False), + (100, True, True, False, False), + (103, True, True, False, False), + (110, True, True, False, False), + (120, True, True, False, False), + (121, True, True, False, False), + ], +) +def test_production_resolver_and_tensor_eligibility_agree( + monkeypatch, capability, triton, tilelang, fusion_expected, fna_expected +): + monkeypatch.setattr(ltx2_platform, "HAS_TRITON", triton) + monkeypatch.setattr(ltx2_platform, "has_tilelang", lambda: tilelang) + monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_cuda", lambda: True) + monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_available", lambda: True) + monkeypatch.setattr( + ltx2_platform.current_omni_platform, + "get_device_capability", + lambda device_id=0: DeviceCapability(major=capability // 10, minor=capability % 10), + ) + # Resolution must be lazy even when FNA is selected. + monkeypatch.setitem(sys.modules, "vllm_omni.diffusion.models.ltx2.ops.fna", None) + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: False) + device = torch.device("cuda:0") + operators = ltx2_platform.resolve_ltx2_vae_operators(device) + assert (operators is not None) is fusion_expected + assert (operators is not None and operators.fna is not None) is fna_expected + if operators is not None: + from vllm_omni.diffusion.models.ltx2.ops import qk_rms_norm, residual_adaln, swiglu + + assert operators.qk_norm_rope is qk_rms_norm.try_qk_rms_norm_scale_rope_3d_exact + assert operators.residual_norm is residual_adaln.try_residual_rms_norm_modulate_exact + assert operators.residual_add is residual_adaln.try_residual_add3_exact + assert operators.swiglu is swiglu.try_swiglu_tiled_exact + + tensor = _TensorProbe(is_cuda=True, device=device) + with torch.inference_mode(): + assert ltx2_platform.is_ltx2_fusion_eligible(tensor) is fusion_expected + assert ltx2_platform.is_ltx2_fna_eligible(tensor) is fna_expected + + +@pytest.mark.parametrize( + ("device", "is_cuda", "available"), + [("cpu", True, True), ("cuda:0", False, True), ("cuda:0", True, False)], +) +def test_resolver_rejects_unsupported_platform_before_capability_lookup(monkeypatch, device, is_cuda, available): + monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_cuda", lambda: is_cuda) + monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_available", lambda: available) + + def unexpected_lookup(_index): + raise AssertionError("Unsupported platforms must not query CUDA capability") + + monkeypatch.setattr(ltx2_platform, "_cuda_compute_capability", unexpected_lookup) + assert ltx2_platform.resolve_ltx2_vae_operators(torch.device(device)) is None + + +@pytest.mark.parametrize(("compiling", "grad_enabled"), [(True, False), (False, True)]) +def test_runtime_guards_bypass_operator_resolution(monkeypatch, compiling, grad_enabled): + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: compiling) + monkeypatch.setattr(torch, "is_grad_enabled", lambda: grad_enabled) + + def unexpected_resolve(_device): + raise AssertionError("Ineligible runtime must not select eager operators") + + monkeypatch.setattr(ltx2_platform, "resolve_ltx2_vae_operators", unexpected_resolve) + tensor = _TensorProbe(is_cuda=True, device=torch.device("cuda:0")) + assert not ltx2_platform.is_ltx2_fusion_eligible(tensor) + assert not ltx2_platform.is_ltx2_fna_eligible(tensor) + + +def test_device_capability_is_cached(monkeypatch): + calls = 0 + monkeypatch.setattr(ltx2_platform, "HAS_TRITON", True) + + def get_device_capability(device_id=0): + nonlocal calls + calls += 1 + return DeviceCapability(major=9, minor=0) + + monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_cuda", lambda: True) + monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_available", lambda: True) + monkeypatch.setattr(ltx2_platform.current_omni_platform, "get_device_capability", get_device_capability) + assert ltx2_platform._cuda_compute_capability(0) == 90 + assert ltx2_platform._cuda_compute_capability(0) == 90 + assert calls == 1 diff --git a/tests/diffusion/models/ltx2/test_ltx2_vae.py b/tests/diffusion/models/ltx2/test_ltx2_vae.py index b9bc52d94..e387eb7a4 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_vae.py +++ b/tests/diffusion/models/ltx2/test_ltx2_vae.py @@ -3,6 +3,8 @@ """Unit tests for LTX video VAE tiling and distributed decode behavior.""" +from collections.abc import Callable +from dataclasses import dataclass from types import SimpleNamespace from typing import Any @@ -35,6 +37,22 @@ def test_ltx_vocoder_deterministic_context(monkeypatch, initial_deterministic, f assert {name: getattr(torch.backends.cudnn, name) for name in settings} == settings +@dataclass(frozen=True) +class _OperatorSetProbe: + fna: Callable[..., torch.Tensor] | None + + +@dataclass +class _AttentionProbe: + kernel_size: tuple[int, int, int] + head_dim: int + to_out: tuple[torch.nn.Module, ...] + projected: torch.Tensor + + def project_qkv(self, _hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return self.projected, self.projected, self.projected + + def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch): import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime @@ -61,9 +79,164 @@ def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch): class TestLTXDiffusionDecoder: + def test_diffusion_decoder_reuses_diffusers_with_scoped_overrides(self, monkeypatch): + from diffusers.models.autoencoders.ltx2_diffusion_decoder import ( + LTX2VideoDiffusionDecoderModel as DiffusersModel, + ) + + from vllm_omni.diffusion.models.ltx2.vae.decoder import ( + LTX2VideoDiffusionDecoderModel, + LTX2VideoVaeDiffusionNABlock, + LTX2VideoVaeNeighborhoodAttention, + LTX2VideoVaeSwiGLU, + ) + + original_init = DiffusersModel.__init__ + original_parameters = {} + + def capture_parameters(self, *args, **kwargs): + original_init(self, *args, **kwargs) + original_parameters.update(self.named_parameters()) + + with torch.device("meta"): + reference = DiffusersModel() + monkeypatch.setattr(DiffusersModel, "__init__", capture_parameters) + model = LTX2VideoDiffusionDecoderModel() + assert isinstance(model, DiffusersModel) + assert all(parameter is original_parameters[name] for name, parameter in model.named_parameters()) + assert {name: tensor.shape for name, tensor in model.state_dict().items()} == { + name: tensor.shape for name, tensor in reference.state_dict().items() + } + + modules = tuple(model.decoder.modules()) + expected_blocks = sum(model.config.decoder_stage_depths) + assert sum(isinstance(module, LTX2VideoVaeNeighborhoodAttention) for module in modules) == expected_blocks + assert sum(isinstance(module, LTX2VideoVaeSwiGLU) for module in modules) == expected_blocks + assert ( + sum(isinstance(module, LTX2VideoVaeDiffusionNABlock) for module in modules) + == model.config.decoder_stage_depths[-1] + ) + + assert type(model.decoder).__module__.startswith("vllm_omni.") + assert model.tiled_decode.__func__ is DiffusersModel.tiled_decode + assert type(reference.decoder) is type(model.decoder).__bases__[0] + + def test_natten_processor_routes_stage5_through_tilelang_fna(self, monkeypatch): + from vllm_omni.diffusion.models.ltx2.vae import decoder as diffvae_modules + + processor = object.__new__(diffvae_modules.LTX2VideoVaeNeighborhoodNattenProcessor) + processor.backend = None + fallback_calls = [] + processor._na3d = lambda *args, **kwargs: fallback_calls.append((args, kwargs)) + + hidden_states = torch.randn(1, 11, 11, 11, 256, dtype=torch.bfloat16) + projected = hidden_states.view(1, 11, 11, 11, 4, 64) + + class FakeAttention: + kernel_size = (11, 11, 11) + head_dim = 64 + to_out = (torch.nn.Identity(),) + + @staticmethod + def project_qkv(_hidden_states): + return projected, projected, projected + + def token_permute(value, _tile, *, flip_tiled_dims): + assert flip_tiled_dims + return value, value.shape, None + + def token_unpermute(value, *, token_layout_shape, tile_shape, flip_tiled_dims): + assert token_layout_shape == projected.shape + assert tile_shape == (4, 4, 4) + assert flip_tiled_dims + return value + + fast_path_calls = [] + + def fast_path(query, key, value, *, shape, window): + fast_path_calls.append((query, key, value, shape, window)) + return query + + monkeypatch.setattr(diffvae_modules, "is_ltx2_fna_eligible", lambda _tensor: True) + monkeypatch.setattr( + diffvae_modules, "resolve_ltx2_vae_operators", lambda _device: _OperatorSetProbe(fna=fast_path) + ) + monkeypatch.setattr( + diffvae_modules, + "_load_natten_token_permutation", + lambda: (token_permute, token_unpermute), + ) + + output = processor(FakeAttention(), hidden_states) + + assert torch.equal(output, hidden_states) + assert len(fast_path_calls) == 1 + assert fast_path_calls[0][3:] == ((11, 11, 11), (11, 11, 11)) + assert not fallback_calls + + @pytest.mark.parametrize("reason", ["device", "head_dim", "channels", "dtype", "tilelang_missing", "kernel"]) + def test_natten_processor_keeps_natten_when_fast_path_does_not_apply(self, monkeypatch, reason): + from vllm_omni.diffusion.models.ltx2.vae import decoder as ops + + processor = object.__new__(ops.LTX2VideoVaeNeighborhoodNattenProcessor) + processor.backend = None + channels = 128 if reason == "channels" else 256 + head_dim = 32 if reason == "head_dim" else 64 + x = torch.zeros(1, 11, 11, 11, channels, dtype=torch.float32 if reason == "dtype" else torch.bfloat16) + projected = x.view(1, 11, 11, 11, channels // head_dim, head_dim) + attn = _AttentionProbe( + kernel_size=(3, 5, 5) if reason == "kernel" else (11, 11, 11), + head_dim=head_dim, + to_out=(torch.nn.Identity(),), + projected=projected, + ) + calls = [] + + def natten(q, k, v, **kwargs): + calls.append(kwargs) + return q + + def unexpected_native_load(*_args, **_kwargs): + raise AssertionError("Native FNA must not be loaded for an unsupported configuration") + + processor._na3d = natten + monkeypatch.setattr(ops, "is_ltx2_fna_eligible", lambda _tensor: reason != "device") + monkeypatch.setattr( + ops, + "resolve_ltx2_vae_operators", + lambda _device: _OperatorSetProbe(fna=None if reason == "tilelang_missing" else unexpected_native_load), + ) + assert torch.equal(processor(attn, x), x) + assert calls == [{"kernel_size": attn.kernel_size, "scale": 1.0, "backend": None}] + + def test_diffusion_decoder_noneligible_block_uses_diffusers_forward(self, monkeypatch): + from vllm_omni.diffusion.models.ltx2.vae import decoder as diffvae_modules + + block = object.__new__(diffvae_modules.LTX2VideoVaeDiffusionNABlock) + torch.nn.Module.__init__(block) + expected = torch.randn(1, 1, 1, 1, 1) + calls = [] + + def upstream_forward(_self, hidden_states, latent_context, modulation, block_mask): + assert block_mask is None + calls.append((hidden_states, latent_context, modulation)) + return expected + + monkeypatch.setattr( + diffvae_modules.diffusers_decoder.LTX2VideoVaeDiffusionNABlock, + "forward", + upstream_forward, + ) + hidden_states = torch.randn(1, 1, 1, 1, 1) + latent_context = torch.randn_like(hidden_states) + modulation = (torch.randn_like(hidden_states),) + + assert block.forward(hidden_states, latent_context, modulation) is expected + assert calls == [(hidden_states, latent_context, modulation)] + @pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width"]) def test_distributed_diffusion_decoder_rejects_non_tile_parallel_modes(self, mode): - from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import ( + from vllm_omni.diffusion.models.ltx2.vae.distributed import ( DistributedLTX2VideoDiffusionDecoderModel, ) @@ -75,7 +248,7 @@ class TestLTXDiffusionDecoder: model.set_parallel_size(2, mode=mode) def test_distributed_diffusion_decoder_accepts_tile_parallel_mode(self): - from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import ( + from vllm_omni.diffusion.models.ltx2.vae.distributed import ( DistributedLTX2VideoDiffusionDecoderModel, ) @@ -91,7 +264,7 @@ class TestLTXDiffusionDecoder: assert calls == [(4, "tile")] def test_short_clip_keeps_stage5_temporal_context_then_crops_output(self): - from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import ( + from vllm_omni.diffusion.models.ltx2.vae.decoder import ( LTX2VideoDiffusionDecoder3d, LTX2VideoDiffusionDecoderModel, ) @@ -166,7 +339,7 @@ class TestLTXDiffusionDecoder: _ltx2_use_diffusion_decoder(SimpleNamespace(extras={"ltx2_use_conv_vae": "true"}), "2.5") def test_native_diffusion_decoder_conversion_splits_qkv_and_folds_gates(self): - from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import ( + from vllm_omni.diffusion.models.ltx2.vae.decoder import ( convert_ltx25_native_diffusion_decoder_state_dict, ) @@ -213,7 +386,7 @@ class TestLTXDiffusionDecoder: assert not any("type_emb" in key or "coarse" in key or "gate_msa" in key for key in converted) def test_native_diffusion_decoder_conversion_rejects_invalid_fused_qkv(self): - from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import ( + from vllm_omni.diffusion.models.ltx2.vae.decoder import ( convert_ltx25_native_diffusion_decoder_state_dict, ) @@ -411,8 +584,8 @@ class TestLTXDiffusionDecoder: assert output.output[1].numel() == 0 def test_diffusion_decoder_patch_parallel_size_one_uses_native_tiling(self, monkeypatch): - from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel - from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import ( + from vllm_omni.diffusion.models.ltx2.vae.decoder import LTX2VideoDiffusionDecoderModel + from vllm_omni.diffusion.models.ltx2.vae.distributed import ( DistributedLTX2VideoDiffusionDecoderModel, ) @@ -447,7 +620,7 @@ class TestLTXDiffusionDecoder: def test_distributed_diffusion_tiles_preserve_serial_noise_order(self, monkeypatch): from diffusers.utils.torch_utils import randn_tensor - from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import ( + from vllm_omni.diffusion.models.ltx2.vae.distributed import ( DistributedLTX2VideoDiffusionDecoderModel, ) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index 72e63a22f..f1dcc38bb 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py @@ -15,7 +15,6 @@ from typing import TYPE_CHECKING, Any import torch from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler -from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoVaeNeighborhoodNattenProcessor from diffusers.pipelines.ltx2 import LTX2TextConnectors from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder @@ -36,12 +35,6 @@ from vllm_omni.transformers_utils.repo_utils import hf_api if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig -from .ltx2_diffusion_decoder import ( - LTX25_NATIVE_ARTIFACT_REVISION, - LTX25_NATIVE_DIFFUSION_DECODER_FILENAME, - LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID, -) -from .ltx2_diffusion_decoder_distributed import DistributedLTX2VideoDiffusionDecoderModel from .ltx2_request import LTXCheckpointKind, validate_ltx_checkpoint from .ltx2_transformer import ( LTX2VideoTransformer3DModel, @@ -49,6 +42,13 @@ from .ltx2_transformer import ( apply_split_rotary_emb, to_ltx_padding_mask, ) +from .vae.decoder import ( + LTX25_NATIVE_ARTIFACT_REVISION, + LTX25_NATIVE_DIFFUSION_DECODER_FILENAME, + LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID, + LTX2VideoVaeNeighborhoodNattenProcessor, +) +from .vae.distributed import DistributedLTX2VideoDiffusionDecoderModel try: from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE diff --git a/vllm_omni/diffusion/models/ltx2/ops/__init__.py b/vllm_omni/diffusion/models/ltx2/ops/__init__.py new file mode 100644 index 000000000..93f2ff672 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ops/__init__.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Shared operator selection for LTX-2 eager kernels and their tests.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from functools import cache + +import torch +from vllm.triton_utils import HAS_TRITON +from vllm.utils.import_utils import has_tilelang + +from vllm_omni.platforms import current_omni_platform + +# Pointwise fusions are launch/memory-traffic reductions whose first execution +# for every runtime key is checked bit-for-bit; failures use the ordinary model +# math. Both paths remain SM90-only until other architectures are validated. +# The FNA schedule is performance-tuned for SM90 and has no silent fallback. +_CUDA_COMPUTE_CAPABILITIES = frozenset({90}) + + +@dataclass(frozen=True) +class LTX2VAEOperatorSet: + qk_norm_rope: Callable[..., tuple[torch.Tensor, torch.Tensor] | None] + residual_norm: Callable[..., torch.Tensor | None] + residual_add: Callable[..., torch.Tensor | None] + swiglu: Callable[..., torch.Tensor | None] + fna: Callable[..., torch.Tensor] | None + + +def _fna3d_tilelang(*args, **kwargs): + # Selecting an operator set must not import the optional compiler backend. + from .fna import fna3d_tilelang + + return fna3d_tilelang(*args, **kwargs) + + +@cache +def _operator_set(with_fna: bool) -> LTX2VAEOperatorSet: + from .qk_rms_norm import try_qk_rms_norm_scale_rope_3d_exact + from .residual_adaln import try_residual_add3_exact, try_residual_rms_norm_modulate_exact + from .swiglu import try_swiglu_tiled_exact + + return LTX2VAEOperatorSet( + qk_norm_rope=try_qk_rms_norm_scale_rope_3d_exact, + residual_norm=try_residual_rms_norm_modulate_exact, + residual_add=try_residual_add3_exact, + swiglu=try_swiglu_tiled_exact, + fna=_fna3d_tilelang if with_fna else None, + ) + + +@cache +def _cuda_compute_capability(device_index: int) -> int | None: + if not HAS_TRITON or not current_omni_platform.is_cuda() or not current_omni_platform.is_available(): + return None + capability = current_omni_platform.get_device_capability(device_id=device_index) + return capability.to_int() if capability is not None else None + + +def resolve_ltx2_vae_operators(device: torch.device) -> LTX2VAEOperatorSet | None: + """Select the production operators, or return None for unsupported devices.""" + if device.type != "cuda" or not current_omni_platform.is_cuda() or not current_omni_platform.is_available(): + return None + device_index = device.index if device.index is not None else torch.accelerator.current_device_index() + if _cuda_compute_capability(int(device_index)) not in _CUDA_COMPUTE_CAPABILITIES: + return None + return _operator_set(has_tilelang()) + + +def _eligible_operators(tensor: torch.Tensor) -> LTX2VAEOperatorSet | None: + if not tensor.is_cuda or torch.compiler.is_compiling() or torch.is_grad_enabled(): + return None + return resolve_ltx2_vae_operators(tensor.device) + + +def is_ltx2_fusion_eligible(tensor: torch.Tensor) -> bool: + """Allow self-verifying pointwise fusions on validated SM90 devices.""" + + return _eligible_operators(tensor) is not None + + +def is_ltx2_fna_eligible(tensor: torch.Tensor) -> bool: + """Keep the strict, Hopper-tuned TileLang FNA schedule on SM90.""" + + operators = _eligible_operators(tensor) + return operators is not None and operators.fna is not None + + +__all__ = ["LTX2VAEOperatorSet", "resolve_ltx2_vae_operators", "is_ltx2_fna_eligible", "is_ltx2_fusion_eligible"] diff --git a/vllm_omni/diffusion/models/ltx2/ops/_numerics.py b/vllm_omni/diffusion/models/ltx2/ops/_numerics.py new file mode 100644 index 000000000..bdc22b142 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ops/_numerics.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Numerical primitives shared by exact LTX-2 Triton kernels.""" + +from vllm.triton_utils import tl, triton + + +@triton.jit +def round_bf16_to_fp32(value): + """RNE-round FP32 to BF16 precision while retaining an FP32 register.""" + + bits = value.to(tl.int32, bitcast=True) + rounding_bias = 0x7FFF + ((bits >> 16) & 1) + rounded_bits = (bits + rounding_bias) & -65536 + return rounded_bits.to(tl.float32, bitcast=True) + + +@triton.jit +def add_rn_f32(x, y): + return tl.inline_asm_elementwise( + asm="add.rn.f32 $0, $1, $2;", + constraints="=f,f,f", + args=[x, y], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def mul_rn_f32(x, y): + return tl.inline_asm_elementwise( + asm="mul.rn.f32 $0, $1, $2;", + constraints="=f,f,f", + args=[x, y], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def fma_rn_f32(x, y, accumulator): + return tl.inline_asm_elementwise( + asm="fma.rn.f32 $0, $1, $2, $3;", + constraints="=f,f,f,f", + args=[x, y, accumulator], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def rsqrt_approx_f32(x): + return tl.inline_asm_elementwise( + asm="rsqrt.approx.f32 $0, $1;", + constraints="=f,f", + args=[x], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def shfl_down_f32(value, delta: tl.constexpr): + return tl.inline_asm_elementwise( + asm="shfl.sync.down.b32 $0, $1, $2, 0x1f, 0xffffffff;", + constraints="=f,f,n", + args=[value, delta], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +__all__ = [ + "add_rn_f32", + "fma_rn_f32", + "mul_rn_f32", + "round_bf16_to_fp32", + "rsqrt_approx_f32", + "shfl_down_f32", +] diff --git a/vllm_omni/diffusion/models/ltx2/ops/fna.py b/vllm_omni/diffusion/models/ltx2/ops/fna.py new file mode 100644 index 000000000..d42477d9a --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ops/fna.py @@ -0,0 +1,525 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""TileLang FNA for the LTX-2.5 DiffVAE stage-5 schedule. + +This intentionally reuses NATTEN's token permutation. It only replaces the +middle BF16 attention kernel for Q tiles (4, 4, 4), KV tiles (4, 4, 8), +head_dim=64, stride=dilation=1, and scale=1.0. +""" + +import functools +import math +from collections import OrderedDict +from dataclasses import dataclass + +import torch +from vllm.tilelang_utils import T, tilelang + +Q_TILE = (4, 4, 4) +KV_TILE = (4, 4, 8) +Q_TOKENS = math.prod(Q_TILE) +KV_TOKENS = math.prod(KV_TILE) +HEAD_DIM = 64 +HEADS = 4 +LOG2_E = 1.4426950408889634 +KV_BUCKET_LIMITS = (32, 48, 60, 75) +PATTERN_TABLE_SIZE = 16 +KV_TILE_ID_BITS = 15 +SHAPE_CACHE_SIZE = 16 + + +@dataclass(frozen=True) +class Problem: + frames: int + height: int + width: int + window: tuple[int, int, int] + + def __post_init__(self) -> None: + if min(self.frames, self.height, self.width) <= 0: + raise ValueError(f"invalid FNA problem: {self}") + if len(self.window) != 3 or any(x <= 0 for x in self.window): + raise ValueError(f"invalid FNA window: {self.window}") + if any(w > n for w, n in zip(self.window, self.shape)): + raise ValueError(f"window {self.window} exceeds shape {self.shape}") + if math.prod(self.kv_grid) > 1 << KV_TILE_ID_BITS: + raise ValueError(f"KV tile grid exceeds {KV_TILE_ID_BITS}-bit metadata encoding: {self.kv_grid}") + + @property + def shape(self) -> tuple[int, int, int]: + return self.frames, self.height, self.width + + @property + def q_padded(self) -> tuple[int, int, int]: + return tuple(_round_up(n, tile) for n, tile in zip(self.shape, Q_TILE)) + + @property + def kv_padded(self) -> tuple[int, int, int]: + return tuple(_round_up(n, tile) for n, tile in zip(self.shape, KV_TILE)) + + @property + def q_grid(self) -> tuple[int, int, int]: + return tuple(n // tile for n, tile in zip(self.q_padded, Q_TILE)) + + @property + def kv_grid(self) -> tuple[int, int, int]: + return tuple(n // tile for n, tile in zip(self.kv_padded, KV_TILE)) + + @property + def q_length(self) -> int: + return math.prod(self.q_padded) + + @property + def kv_length(self) -> int: + return math.prod(self.kv_padded) + + +def _round_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple + + +def _window_start(index: int, window: int, length: int) -> int: + return min(max(index - window // 2, 0), length - window) + + +@functools.lru_cache(maxsize=32) +def _dimension_pattern_table( + length: int, + q_tile_size: int, + kv_tile_size: int, + window: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Map (Q tile, KV tile) pairs to a small table of packed 1D masks.""" + q_tiles = _round_up(length, q_tile_size) // q_tile_size + kv_tiles = _round_up(length, kv_tile_size) // kv_tile_size + zero = (0,) * q_tile_size + patterns: list[tuple[int, ...]] = [zero] + pattern_ids = {zero: 0} + ids = torch.zeros((q_tiles, kv_tiles), dtype=torch.int64) + + for q_tile in range(q_tiles): + for kv_tile in range(kv_tiles): + rows: list[int] = [] + for q_inner in range(q_tile_size): + q = q_tile * q_tile_size + q_inner + bits = 0 + if q < length: + start = _window_start(q, window, length) + for kv_inner in range(kv_tile_size): + kv = kv_tile * kv_tile_size + kv_inner + if kv < length and start <= kv < start + window: + bits |= 1 << kv_inner + rows.append(bits) + signature = tuple(rows) + pattern_id = pattern_ids.get(signature) + if pattern_id is None: + pattern_id = len(patterns) + pattern_ids[signature] = pattern_id + patterns.append(signature) + ids[q_tile, kv_tile] = pattern_id + + if len(patterns) > PATTERN_TABLE_SIZE: + raise AssertionError(f"4-bit pattern id overflow for length={length}: {len(patterns)}") + packed = [] + for pattern in patterns: + value = sum(bits << (q * kv_tile_size) for q, bits in enumerate(pattern)) + # Store uint32 bit patterns in an int32 tensor. Arithmetic right shift + # followed by `& 1` still extracts the desired bit when bit 31 is set. + packed.append(value if value < (1 << 31) else value - (1 << 32)) + # The four-bit pattern ids address at most 16 entries. Pad to that fixed + # size so T/H/W-specific table lengths cannot specialize the GPU kernel. + packed.extend([0] * (PATTERN_TABLE_SIZE - len(packed))) + return ids, torch.tensor(packed, dtype=torch.int32) + + +@functools.lru_cache(maxsize=SHAPE_CACHE_SIZE) +def build_kv_tile_map(problem: Problem) -> tuple[torch.Tensor, torch.Tensor]: + """Return packed KV tile ids and valid counts for each permuted Q tile.""" + q_grid_t, q_grid_h, q_grid_w = problem.q_grid + kv_grid_t, kv_grid_h, _ = problem.kv_grid + rows: list[list[int]] = [] + + # Token permutation flattens outer tile modes as (W, H, T), with T + # changing fastest. Keep KV tiles in the same increasing linear order. + for q_w_tile in range(q_grid_w): + for q_h_tile in range(q_grid_h): + for q_t_tile in range(q_grid_t): + ranges: list[range] = [] + for q_tile_idx, q_tile_size, kv_tile_size, window, length in zip( + (q_t_tile, q_h_tile, q_w_tile), + Q_TILE, + KV_TILE, + problem.window, + problem.shape, + ): + first_q = q_tile_idx * q_tile_size + last_q = min(first_q + q_tile_size, length) - 1 + first_k = _window_start(first_q, window, length) + last_k = _window_start(last_q, window, length) + window - 1 + ranges.append(range(first_k // kv_tile_size, last_k // kv_tile_size + 1)) + + t_tiles, h_tiles, w_tiles = ranges + row = [ + (kv_w * kv_grid_h + kv_h) * kv_grid_t + kv_t + for kv_w in w_tiles + for kv_h in h_tiles + for kv_t in t_tiles + ] + if not row: + raise AssertionError("every Q tile must overlap at least one KV tile") + rows.append(row) + + max_tiles = max(map(len, rows)) + # Duplicate a valid tile in padding slots. The kernel masks those slots + # with `counts`; avoiding -1 keeps every indirect global load in bounds. + indices = torch.tensor( + [row + [row[0]] * (max_tiles - len(row)) for row in rows], + dtype=torch.int32, + ) + counts = torch.tensor([len(row) for row in rows], dtype=torch.int32) + return indices, counts + + +@functools.lru_cache(maxsize=SHAPE_CACHE_SIZE) +def build_kv_tile_metadata( + problem: Problem, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Pack KV tile id and three separable mask pattern ids into int32.""" + indices, counts = build_kv_tile_map(problem) + q_grid_t, q_grid_h, _ = problem.q_grid + kv_grid_t, kv_grid_h, _ = problem.kv_grid + t_ids, t_patterns = _dimension_pattern_table(problem.frames, Q_TILE[0], KV_TILE[0], problem.window[0]) + h_ids, h_patterns = _dimension_pattern_table(problem.height, Q_TILE[1], KV_TILE[1], problem.window[1]) + w_ids, w_patterns = _dimension_pattern_table(problem.width, Q_TILE[2], KV_TILE[2], problem.window[2]) + + q_blocks = torch.arange(indices.shape[0], dtype=torch.int64)[:, None] + q_t = q_blocks % q_grid_t + q_h = (q_blocks // q_grid_t) % q_grid_h + q_w = q_blocks // (q_grid_t * q_grid_h) + kv = indices.to(torch.int64) + kv_t = kv % kv_grid_t + kv_h = (kv // kv_grid_t) % kv_grid_h + kv_w = kv // (kv_grid_t * kv_grid_h) + t_pattern_id = t_ids[q_t, kv_t] + h_pattern_id = h_ids[q_h, kv_h] + w_pattern_id = w_ids[q_w, kv_w] + metadata = kv | (t_pattern_id << 15) | (h_pattern_id << 19) | (w_pattern_id << 23) + active = torch.arange(indices.shape[1])[None, :] < counts[:, None] + metadata = torch.where(active, metadata, 0).to(torch.int32) + return metadata, t_patterns, h_patterns, w_patterns + + +@functools.lru_cache(maxsize=SHAPE_CACHE_SIZE) +def build_kv_tile_buckets( + problem: Problem, +) -> tuple[tuple[torch.Tensor, torch.Tensor, int], ...]: + """Group Q tiles by an upper bound on their useful KV tile count.""" + metadata, *_ = build_kv_tile_metadata(problem) + _, counts = build_kv_tile_map(problem) + maximum = int(counts.max()) + if maximum > KV_BUCKET_LIMITS[-1]: + raise ValueError( + f"FNA neighborhood needs {maximum} KV tiles, exceeding the supported ceiling {KV_BUCKET_LIMITS[-1]}" + ) + buckets = [] + lower = 0 + for limit in KV_BUCKET_LIMITS: + q_tile_ids = torch.nonzero((counts > lower) & (counts <= limit), as_tuple=False).flatten() + if q_tile_ids.numel(): + # Small grids can have fewer columns than the first bucket limit. + # Zero metadata encodes an all-masked tile, preserving the fixed ABI. + bucket_metadata = torch.zeros((q_tile_ids.numel(), limit), dtype=torch.int32) + columns = min(limit, metadata.shape[1]) + bucket_metadata[:, :columns] = metadata[q_tile_ids, :columns] + buckets.append( + ( + q_tile_ids.to(torch.int32), + bucket_metadata, + limit, + ) + ) + lower = limit + if limit >= maximum: + break + return tuple(buckets) + + +@functools.lru_cache(maxsize=32) +def compile_kernel( + max_kv_tiles: int, + *, + target: str = "auto", +): + """Compile one T/H/W-dynamic kernel for a fixed KV-count ceiling.""" + if max_kv_tiles not in KV_BUCKET_LIMITS: + raise ValueError(f"unsupported KV bucket ceiling: {max_kv_tiles}") + + q_length = T.dynamic("q_length") + kv_length = T.dynamic("kv_length") + q_tile_count = T.dynamic("q_tile_count") + + # Batch is fixed to one by the public guard. Excluding it from the + # kernel ABI keeps every physical stride static even when token lengths + # vary, preserving vectorized global/shared-memory copies. + q_shape = [q_length, HEADS, HEAD_DIM] + kv_shape = [kv_length, HEADS, HEAD_DIM] + q_tile_id_shape = [q_tile_count] + map_shape = [q_tile_count, max_kv_tiles] + pattern_shape = [PATTERN_TABLE_SIZE] + dtype = T.bfloat16 + accum_dtype = T.float32 + + @T.prim_func + def main( + query: T.Tensor(q_shape, dtype), + key: T.Tensor(kv_shape, dtype), + value: T.Tensor(kv_shape, dtype), + q_tile_ids: T.Tensor(q_tile_id_shape, T.int32), + kv_tile_metadata: T.Tensor(map_shape, T.int32), + t_mask_patterns: T.Tensor(pattern_shape, T.int32), + h_mask_patterns: T.Tensor(pattern_shape, T.int32), + w_mask_patterns: T.Tensor(pattern_shape, T.int32), + output: T.Tensor(q_shape, dtype), + ): + with T.Kernel(q_tile_count, HEADS, threads=128) as (bx, by): + q_shared = T.alloc_shared([Q_TOKENS, HEAD_DIM], dtype) + k_shared = T.alloc_shared([KV_TOKENS, HEAD_DIM], dtype) + v_shared = T.alloc_shared([KV_TOKENS, HEAD_DIM], dtype) + o_shared = T.alloc_shared([Q_TOKENS, HEAD_DIM], dtype) + scores = T.alloc_fragment([Q_TOKENS, KV_TOKENS], accum_dtype) + probs = T.alloc_fragment([Q_TOKENS, KV_TOKENS], dtype) + acc_o = T.alloc_fragment([Q_TOKENS, HEAD_DIM], accum_dtype) + row_max = T.alloc_fragment([Q_TOKENS], accum_dtype) + row_max_prev = T.alloc_fragment([Q_TOKENS], accum_dtype) + row_scale = T.alloc_fragment([Q_TOKENS], accum_dtype) + row_sum = T.alloc_fragment([Q_TOKENS], accum_dtype) + row_denom = T.alloc_fragment([Q_TOKENS], accum_dtype) + + q_block = q_tile_ids[bx] + # CPU metadata only emits padded Q/KV tiles that are fully backed + # by the permuted tensors. Expose those invariants so dynamic + # lengths do not add per-copy boundary predicates. + T.assume(q_block >= 0) + T.assume((q_block + 1) * Q_TOKENS <= q_length) + T.copy( + query[ + q_block * Q_TOKENS : (q_block + 1) * Q_TOKENS, + by, + :, + ], + q_shared, + ) + T.fill(acc_o, 0) + T.fill(row_denom, 0) + T.fill(row_max, -T.infinity(accum_dtype)) + + for block_iter in T.Pipelined(max_kv_tiles, num_stages=1): + metadata = kv_tile_metadata[bx, block_iter] + kv_block = metadata & 0x7FFF + T.assume((kv_block + 1) * KV_TOKENS <= kv_length) + t_mask = t_mask_patterns[(metadata >> 15) & 0xF] + h_mask = h_mask_patterns[(metadata >> 19) & 0xF] + w_mask = w_mask_patterns[(metadata >> 23) & 0xF] + T.copy( + key[ + kv_block * KV_TOKENS : (kv_block + 1) * KV_TOKENS, + by, + :, + ], + k_shared, + prefer_instruction="cp_async", + ) + + for i, j in T.Parallel(Q_TOKENS, KV_TOKENS): + q_t = i % Q_TILE[0] + q_h = (i // Q_TILE[0]) % Q_TILE[1] + q_w = i // (Q_TILE[0] * Q_TILE[1]) + k_t = j % KV_TILE[0] + k_h = (j // KV_TILE[0]) % KV_TILE[1] + k_w = j // (KV_TILE[0] * KV_TILE[1]) + valid = ( + (((t_mask >> (q_t * KV_TILE[0] + k_t)) & 1) != 0) + & (((h_mask >> (q_h * KV_TILE[1] + k_h)) & 1) != 0) + & (((w_mask >> (q_w * KV_TILE[2] + k_w)) & 1) != 0) + ) + scores[i, j] = T.if_then_else(valid, 0, -T.infinity(accum_dtype)) + + T.gemm( + q_shared, + k_shared, + scores, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(row_max, row_max_prev) + T.fill(row_max, -T.infinity(accum_dtype)) + T.reduce_max(scores, row_max, dim=1, clear=False) + # A KV tile is selected for the union of all 64 queries in + # this Q tile. It can therefore be fully masked for one + # individual query. Preserve that row's previous online + # softmax state instead of evaluating (-inf) - (-inf). + T.copy(row_max, row_sum) + for i in T.Parallel(Q_TOKENS): + has_score = row_sum[i] > -T.infinity(accum_dtype) + merged_max = T.max(row_sum[i], row_max_prev[i]) + row_max[i] = T.if_then_else(has_score, merged_max, row_max_prev[i]) + row_scale[i] = T.if_then_else( + has_score, + T.exp2((row_max_prev[i] - merged_max) * LOG2_E), + 1, + ) + # Match Hopper NATTEN's two separately rounded FP32 + # products before exp2. + row_sum[i] = T.if_then_else( + has_score, + merged_max * LOG2_E, + -T.infinity(accum_dtype), + ) + for i, j in T.Parallel(Q_TOKENS, KV_TOKENS): + scores[i, j] = T.if_then_else( + row_sum[i] > -T.infinity(accum_dtype), + T.exp2(scores[i, j] * LOG2_E - row_sum[i]), + 0, + ) + T.reduce_sum(scores, row_sum, dim=1) + for i in T.Parallel(Q_TOKENS): + row_denom[i] = row_denom[i] * row_scale[i] + row_sum[i] + T.copy(scores, probs) + for i, j in T.Parallel(Q_TOKENS, HEAD_DIM): + acc_o[i, j] *= row_scale[i] + + T.copy( + value[ + kv_block * KV_TOKENS : (kv_block + 1) * KV_TOKENS, + by, + :, + ], + v_shared, + prefer_instruction="cp_async", + ) + T.gemm(probs, v_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) + + # Match NATTEN's epilogue: compute one reciprocal per query row, + # then multiply every output accumulator by that shared alpha. + for i in T.Parallel(Q_TOKENS): + row_scale[i] = T.if_then_else( + row_denom[i] > 0, + T.call_extern("float32", "__frcp_rn", row_denom[i]), + 0, + ) + for i, j in T.Parallel(Q_TOKENS, HEAD_DIM): + acc_o[i, j] *= row_scale[i] + T.copy(acc_o, o_shared) + T.copy( + o_shared, + output[ + q_block * Q_TOKENS : (q_block + 1) * Q_TOKENS, + by, + :, + ], + ) + + return tilelang.compile( + main, + out_idx=None, + target="auto" if target == "auto" else {"kind": "cuda", "arch": target}, + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + # Keeping all work in the 128-thread consumer warpgroup avoids a + # separate producer warpgroup. The generated 240-register + # consumer can then admit two CTAs per SM instead of one. + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + }, + ) + + +_DEVICE_MAP_CACHE: OrderedDict[tuple[Problem, torch.device], tuple[object, ...]] = OrderedDict() + + +def _fna3d_tilelang( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + shape: tuple[int, int, int], + window: tuple[int, int, int], +) -> torch.Tensor: + problem = Problem(*shape, window=window) + if query.shape[1] != problem.q_length or key.shape[1] != problem.kv_length: + raise ValueError(f"permuted lengths do not match {problem}: Q={query.shape[1]}, KV={key.shape[1]}") + cache_key = (problem, query.device) + device_inputs = _DEVICE_MAP_CACHE.get(cache_key) + if device_inputs is None: + _, t_patterns, h_patterns, w_patterns = build_kv_tile_metadata(problem) + pattern_tables = tuple( + tensor.to(query.device, non_blocking=True) for tensor in (t_patterns, h_patterns, w_patterns) + ) + buckets = tuple( + ( + q_tile_ids.to(query.device, non_blocking=True), + metadata.to(query.device, non_blocking=True), + limit, + ) + for q_tile_ids, metadata, limit in build_kv_tile_buckets(problem) + ) + device_inputs = (pattern_tables, buckets) + _DEVICE_MAP_CACHE[cache_key] = device_inputs + if len(_DEVICE_MAP_CACHE) > SHAPE_CACHE_SIZE: + _DEVICE_MAP_CACHE.popitem(last=False) + else: + _DEVICE_MAP_CACHE.move_to_end(cache_key) + + pattern_tables, buckets = device_inputs + output = torch.empty_like(query) + for q_tile_ids, metadata, limit in buckets: + from tilelang.backend.target import determine_target + + target = str(determine_target("auto", return_object=True).attrs["arch"]) + kernel = compile_kernel(limit, target=target) + for batch in range(query.shape[0]): + kernel(query[batch], key[batch], value[batch], q_tile_ids, metadata, *pattern_tables, output[batch]) + return output + + +def fna3d_tilelang( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + shape: tuple[int, int, int], + window: tuple[int, int, int], +) -> torch.Tensor: + """Run native stage-5 FNA, or raise for unsupported inputs/launch failures. + + Q and KV use the documented token-tiled layout. Batch members share spatial + metadata; compilation is cached by target architecture and KV-count bucket. + """ + if not ( + query.is_cuda + and query.dtype is torch.bfloat16 + and query.ndim == key.ndim == value.ndim == 4 + and query.shape[0] > 0 + and query.shape[-2:] == (HEADS, HEAD_DIM) + and query.numel() > 0 + and query.is_contiguous() + and key.device == value.device == query.device + and key.dtype is value.dtype is query.dtype + and key.is_contiguous() + and value.is_contiguous() + and value.shape == key.shape + and key.shape[0] == query.shape[0] + and key.shape[-2:] == query.shape[-2:] + and len(shape) == 3 + and tuple(window) == (11, 11, 11) + ): + raise ValueError( + "LTX DiffVAE FNA requires contiguous CUDA BF16 [B, tokens, 4, 64] Q/K/V and an 11x11x11 window." + ) + if torch.is_grad_enabled() and any(x.requires_grad for x in (query, key, value)): + raise RuntimeError("LTX DiffVAE FNA is inference-only.") + with torch.accelerator.device_index(query.device.index): + return _fna3d_tilelang(query, key, value, shape=tuple(shape), window=tuple(window)) + + +__all__ = ["fna3d_tilelang"] diff --git a/vllm_omni/diffusion/models/ltx2/ops/qk_rms_norm.py b/vllm_omni/diffusion/models/ltx2/ops/qk_rms_norm.py new file mode 100644 index 000000000..be89aeca2 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ops/qk_rms_norm.py @@ -0,0 +1,377 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Bit-exact multi-row Q/K RMSNorm for the LTX diffusion VAE decoder.""" + +from __future__ import annotations + +import logging +import math + +import torch +import torch.nn.functional as F +from vllm.triton_utils import tl, triton + +from . import is_ltx2_fusion_eligible +from ._numerics import ( + add_rn_f32, + fma_rn_f32, + mul_rn_f32, + round_bf16_to_fp32, + rsqrt_approx_f32, + shfl_down_f32, +) +from .rope_tables import get_rope_tables + +_HEAD_DIM = 64 +_ROWS_PER_PROGRAM = 2 +_VERIFY_ROWS = 256 +_FAILED_ROPE_KEYS: set[tuple[int | None, float, float, tuple[int, int, int], float]] = set() +_VERIFIED_ROPE_KEYS: set[tuple[int | None, float, float, tuple[int, int, int], float]] = set() + +logger = logging.getLogger(__name__) + + +@triton.jit +def _shuffle_lane_zero(value): + return tl.inline_asm_elementwise( + asm="shfl.sync.idx.b32 $0, $1, 0, 0x1f, 0xffffffff;", + constraints="=f,f", + args=[value], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _warp_sum(value): + value = add_rn_f32(value, shfl_down_f32(value, 16)) + value = add_rn_f32(value, shfl_down_f32(value, 8)) + value = add_rn_f32(value, shfl_down_f32(value, 4)) + value = add_rn_f32(value, shfl_down_f32(value, 2)) + value = add_rn_f32(value, shfl_down_f32(value, 1)) + return _shuffle_lane_zero(value) + + +@triton.jit +def _qk_rms_norm_scale_rope_3d_kernel( + query_output_ptr, + key_output_ptr, + query_ptr, + key_ptr, + query_weight_ptr, + key_weight_ptr, + cos_t_ptr, + sin_t_ptr, + cos_h_ptr, + sin_h_ptr, + cos_w_ptr, + sin_w_ptr, + rows, + tokens_per_batch, + height, + width, + heads, + eps, + query_scale, + rows_per_program: tl.constexpr, + head_dim: tl.constexpr, + pairs_t: tl.constexpr, + pairs_h: tl.constexpr, +): + threads = tl.arange(0, rows_per_program * 32) + local_row = threads // 32 + lane = threads % 32 + row = tl.program_id(0).to(tl.int64) * rows_per_program + local_row + valid_row = row < rows + reduction_valid = valid_row & (lane < 16) + reduction_base = row * head_dim + lane * 4 + + query_accumulator = tl.zeros((rows_per_program * 32,), dtype=tl.float32) + key_accumulator = tl.zeros((rows_per_program * 32,), dtype=tl.float32) + for element in tl.static_range(4): + query = tl.load(query_ptr + reduction_base + element, mask=reduction_valid, other=0.0).to(tl.float32) + key = tl.load(key_ptr + reduction_base + element, mask=reduction_valid, other=0.0).to(tl.float32) + query_accumulator = fma_rn_f32(query, query, query_accumulator) + key_accumulator = fma_rn_f32(key, key, key_accumulator) + query_rstd = rsqrt_approx_f32(_warp_sum(query_accumulator) / head_dim + eps) + key_rstd = rsqrt_approx_f32(_warp_sum(key_accumulator) / head_dim + eps) + + tokens = (row // heads) % tokens_per_batch + pos_w = tokens % width + pos_h = (tokens // width) % height + pos_t = tokens // (height * width) + pairs_per_head: tl.constexpr = head_dim // 2 + pairs_w: tl.constexpr = pairs_per_head - pairs_t - pairs_h + is_t = lane < pairs_t + is_h = (lane >= pairs_t) & (lane < pairs_t + pairs_h) + is_w = ~(is_t | is_h) + pair_t = lane + pair_h = lane - pairs_t + pair_w = lane - pairs_t - pairs_h + cos_t = tl.load(cos_t_ptr + pos_t * pairs_t + pair_t, mask=valid_row & is_t, other=0.0) + sin_t = tl.load(sin_t_ptr + pos_t * pairs_t + pair_t, mask=valid_row & is_t, other=0.0) + cos_h = tl.load(cos_h_ptr + pos_h * pairs_h + pair_h, mask=valid_row & is_h, other=0.0) + sin_h = tl.load(sin_h_ptr + pos_h * pairs_h + pair_h, mask=valid_row & is_h, other=0.0) + cos_w = tl.load(cos_w_ptr + pos_w * pairs_w + pair_w, mask=valid_row & is_w, other=0.0) + sin_w = tl.load(sin_w_ptr + pos_w * pairs_w + pair_w, mask=valid_row & is_w, other=0.0) + cos = cos_t + cos_h + cos_w + sin = sin_t + sin_h + sin_w + + even_column = lane * 2 + odd_column = even_column + 1 + even_offset = row * head_dim + even_column + odd_offset = row * head_dim + odd_column + query_even = tl.load(query_ptr + even_offset, mask=valid_row).to(tl.float32) + query_odd = tl.load(query_ptr + odd_offset, mask=valid_row).to(tl.float32) + key_even = tl.load(key_ptr + even_offset, mask=valid_row).to(tl.float32) + key_odd = tl.load(key_ptr + odd_offset, mask=valid_row).to(tl.float32) + query_even_weight = tl.load(query_weight_ptr + even_column).to(tl.float32) + query_odd_weight = tl.load(query_weight_ptr + odd_column).to(tl.float32) + key_even_weight = tl.load(key_weight_ptr + even_column).to(tl.float32) + key_odd_weight = tl.load(key_weight_ptr + odd_column).to(tl.float32) + + query_even = round_bf16_to_fp32(mul_rn_f32(query_even_weight, mul_rn_f32(query_rstd, query_even))) + query_odd = round_bf16_to_fp32(mul_rn_f32(query_odd_weight, mul_rn_f32(query_rstd, query_odd))) + query_even = round_bf16_to_fp32(mul_rn_f32(query_even, query_scale)) + query_odd = round_bf16_to_fp32(mul_rn_f32(query_odd, query_scale)) + key_even = round_bf16_to_fp32(mul_rn_f32(key_even_weight, mul_rn_f32(key_rstd, key_even))) + key_odd = round_bf16_to_fp32(mul_rn_f32(key_odd_weight, mul_rn_f32(key_rstd, key_odd))) + + tl.store( + query_output_ptr + even_offset, + add_rn_f32(mul_rn_f32(query_even, cos), -mul_rn_f32(query_odd, sin)), + mask=valid_row, + ) + tl.store( + query_output_ptr + odd_offset, + add_rn_f32(mul_rn_f32(query_even, sin), mul_rn_f32(query_odd, cos)), + mask=valid_row, + ) + tl.store( + key_output_ptr + even_offset, + add_rn_f32(mul_rn_f32(key_even, cos), -mul_rn_f32(key_odd, sin)), + mask=valid_row, + ) + tl.store( + key_output_ptr + odd_offset, + add_rn_f32(mul_rn_f32(key_even, sin), mul_rn_f32(key_odd, cos)), + mask=valid_row, + ) + + +def _weight_matches(value: torch.Tensor, weight: torch.Tensor) -> bool: + return ( + weight.is_cuda + and weight.device == value.device + and weight.dtype is torch.bfloat16 + and weight.shape == (_HEAD_DIM,) + and weight.is_contiguous() + ) + + +def _supported_inputs( + query: torch.Tensor, + key: torch.Tensor, + query_weight: torch.Tensor, + key_weight: torch.Tensor, + eps: float, + query_scale: float, +) -> bool: + return ( + is_ltx2_fusion_eligible(query) + and query.dtype is torch.bfloat16 + and query.ndim == 6 + and query.shape[-1] == _HEAD_DIM + and query.numel() > 0 + and query.is_contiguous() + and key.is_cuda + and key.device == query.device + and key.dtype is query.dtype + and key.shape == query.shape + and key.is_contiguous() + and _weight_matches(query, query_weight) + and _weight_matches(query, key_weight) + and math.isfinite(eps) + and eps > 0 + and math.isfinite(query_scale) + ) + + +def _combined_supported_inputs( + query: torch.Tensor, + key: torch.Tensor, + query_weight: torch.Tensor, + key_weight: torch.Tensor, + eps: float, + query_scale: float, + dim_split: tuple[int, int, int], + base: float, +) -> bool: + return ( + _supported_inputs(query, key, query_weight, key_weight, eps, query_scale) + and len(dim_split) == 3 + and sum(dim_split) == _HEAD_DIM + and all(dim > 0 and dim % 2 == 0 for dim in dim_split) + and math.isfinite(base) + and base > 0 + ) + + +def _launch_combined( + query: torch.Tensor, + key: torch.Tensor, + query_weight: torch.Tensor, + key_weight: torch.Tensor, + eps: float, + query_scale: float, + dim_split: tuple[int, int, int], + base: float, +) -> tuple[torch.Tensor, torch.Tensor]: + query_output = torch.empty_like(query) + key_output = torch.empty_like(key) + rows = query.numel() // _HEAD_DIM + tables = get_rope_tables(query, dim_split, base) + cos_t, sin_t = tables[0] + cos_h, sin_h = tables[1] + cos_w, sin_w = tables[2] + with torch.accelerator.device_index(query.device.index): + _qk_rms_norm_scale_rope_3d_kernel[(triton.cdiv(rows, _ROWS_PER_PROGRAM),)]( + query_output, + key_output, + query, + key, + query_weight, + key_weight, + cos_t, + sin_t, + cos_h, + sin_h, + cos_w, + sin_w, + rows, + query.shape[1] * query.shape[2] * query.shape[3], + query.shape[2], + query.shape[3], + query.shape[4], + eps, + query_scale, + rows_per_program=_ROWS_PER_PROGRAM, + head_dim=_HEAD_DIM, + pairs_t=dim_split[0] // 2, + pairs_h=dim_split[1] // 2, + num_warps=_ROWS_PER_PROGRAM, + ) + return query_output, key_output + + +def _verify_combined_prefix( + output: tuple[torch.Tensor, torch.Tensor], + query: torch.Tensor, + key: torch.Tensor, + query_weight: torch.Tensor, + key_weight: torch.Tensor, + eps: float, + query_scale: float, + dim_split: tuple[int, int, int], + base: float, +) -> bool: + rows = min(query.numel() // _HEAD_DIM, _VERIFY_ROWS) + query_rows = query.reshape(-1, _HEAD_DIM)[:rows] + key_rows = key.reshape(-1, _HEAD_DIM)[:rows] + query_rows = F.rms_norm(query_rows, (_HEAD_DIM,), query_weight, eps) * query_scale + key_rows = F.rms_norm(key_rows, (_HEAD_DIM,), key_weight, eps) + tables = get_rope_tables(query, dim_split, base) + row_indices = torch.arange(rows, device=query.device) + tokens_per_batch = query.shape[1] * query.shape[2] * query.shape[3] + tokens = (row_indices // query.shape[4]) % tokens_per_batch + positions = ( + tokens // (query.shape[2] * query.shape[3]), + (tokens // query.shape[3]) % query.shape[2], + tokens % query.shape[3], + ) + cos = torch.cat([axis_tables[0][position] for axis_tables, position in zip(tables, positions, strict=True)], -1) + sin = torch.cat([axis_tables[1][position] for axis_tables, position in zip(tables, positions, strict=True)], -1) + + def rotate(value: torch.Tensor) -> torch.Tensor: + pairs = value.reshape(rows, _HEAD_DIM // 2, 2) + even = pairs[..., 0].float() + odd = pairs[..., 1].float() + return ( + torch.stack((even * cos - odd * sin, even * sin + odd * cos), -1).reshape(rows, _HEAD_DIM).to(value.dtype) + ) + + reference = rotate(query_rows), rotate(key_rows) + actual = tuple(value.reshape(-1, _HEAD_DIM)[:rows] for value in output) + return all(torch.equal(a, b) for a, b in zip(actual, reference, strict=True)) + + +def try_qk_rms_norm_scale_rope_3d_exact( + query: torch.Tensor, + key: torch.Tensor, + query_weight: torch.Tensor, + key_weight: torch.Tensor, + eps: float, + query_scale: float, + dim_split: tuple[int, int, int], + base: float, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Normalize Q/K, scale Q, and apply paired 3D RoPE in one kernel.""" + + if not _combined_supported_inputs( + query, + key, + query_weight, + key_weight, + eps, + query_scale, + dim_split, + base, + ): + return None + runtime_key = (query.device.index, float(eps), float(query_scale), dim_split, float(base)) + if runtime_key in _FAILED_ROPE_KEYS: + return None + try: + output = _launch_combined( + query, + key, + query_weight, + key_weight, + eps, + query_scale, + dim_split, + base, + ) + if runtime_key not in _VERIFIED_ROPE_KEYS: + if not _verify_combined_prefix( + output, + query, + key, + query_weight, + key_weight, + eps, + query_scale, + dim_split, + base, + ): + _FAILED_ROPE_KEYS.add(runtime_key) + logger.warning( + "Disabling LTX DiffVAE Q/K RMSNorm+3D RoPE fusion on %s after a bit-exactness mismatch", + query.device, + ) + return None + _VERIFIED_ROPE_KEYS.add(runtime_key) + except Exception as exc: # noqa: BLE001 - fail closed after optimized-path failure + _FAILED_ROPE_KEYS.add(runtime_key) + logger.warning( + "Disabling LTX DiffVAE Q/K RMSNorm+3D RoPE fusion on %s after failure: %s", + query.device, + exc, + ) + return None + return output + + +__all__ = ["try_qk_rms_norm_scale_rope_3d_exact"] diff --git a/vllm_omni/diffusion/models/ltx2/ops/residual_adaln.py b/vllm_omni/diffusion/models/ltx2/ops/residual_adaln.py new file mode 100644 index 000000000..16a4aefe6 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ops/residual_adaln.py @@ -0,0 +1,489 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Exact LTX DiffVAE residual, RMSNorm, and AdaLN Triton fusions.""" + +from __future__ import annotations + +import logging +import math + +import torch +import torch.nn.functional as F +from vllm.triton_utils import tl, triton + +from . import is_ltx2_fusion_eligible +from ._numerics import ( + add_rn_f32, + fma_rn_f32, + mul_rn_f32, + round_bf16_to_fp32, + rsqrt_approx_f32, + shfl_down_f32, +) + +_HIDDEN_SIZE = 256 +_POINTWISE_BLOCK = 1024 +_VERIFY_ROWS = 256 +_VERIFY_ELEMENTS = 1 << 20 +_FAILED_ADALN_KEYS: set[tuple[int | None, int, float]] = set() +_VERIFIED_ADALN_KEYS: set[tuple[int | None, int, float]] = set() +_FAILED_ADD_DEVICES: set[int | None] = set() +_VERIFIED_ADD_DEVICES: set[int | None] = set() + +logger = logging.getLogger(__name__) + + +@triton.jit +def _two_warp_sum(values): + values = add_rn_f32(values, shfl_down_f32(values, 16)) + values = add_rn_f32(values, shfl_down_f32(values, 8)) + values = add_rn_f32(values, shfl_down_f32(values, 4)) + values = add_rn_f32(values, shfl_down_f32(values, 2)) + values = add_rn_f32(values, shfl_down_f32(values, 1)) + threads = tl.arange(0, 64) + warp_0 = tl.sum(tl.where(threads == 0, values, 0.0)) + warp_1 = tl.sum(tl.where(threads == 32, values, 0.0)) + return add_rn_f32(warp_0, warp_1) + + +@triton.jit +def _load_ordered_residual( + x_ptr, + residual_a_ptr, + residual_b_ptr, + offset, + valid, + residual_count: tl.constexpr, +): + value = add_rn_f32( + tl.load(x_ptr + offset, mask=valid, other=0.0).to(tl.float32), + tl.load(residual_a_ptr + offset, mask=valid, other=0.0).to(tl.float32), + ) + value = round_bf16_to_fp32(value) + if residual_count == 2: + value = add_rn_f32( + value, + tl.load(residual_b_ptr + offset, mask=valid, other=0.0).to(tl.float32), + ) + value = round_bf16_to_fp32(value) + return value + + +@triton.jit +def _store_modulated( + output_ptr, + norm_weight_ptr, + scale_ptr, + shift_ptr, + value, + offset, + column, + modulation_base, + reciprocal_rms, + valid, +): + normalized = mul_rn_f32(value, reciprocal_rms) + norm_weight = tl.load(norm_weight_ptr + column, mask=valid, other=0.0).to(tl.float32) + normalized = round_bf16_to_fp32(mul_rn_f32(normalized, norm_weight)) + scale = tl.load(scale_ptr + modulation_base + column, mask=valid, other=0.0).to(tl.float32) + shift = tl.load(shift_ptr + modulation_base + column, mask=valid, other=0.0).to(tl.float32) + one_plus_scale = round_bf16_to_fp32(add_rn_f32(1.0, scale)) + product = round_bf16_to_fp32(mul_rn_f32(normalized, one_plus_scale)) + tl.store(output_ptr + offset, add_rn_f32(product, shift), mask=valid) + + +@triton.jit +def _residual_rms_norm_modulate_kernel( + output_ptr, + x_ptr, + residual_a_ptr, + residual_b_ptr, + norm_weight_ptr, + scale_ptr, + shift_ptr, + rows_per_batch, + eps, + residual_count: tl.constexpr, + hidden_size: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + batch = row // rows_per_batch + threads = tl.arange(0, 64) + row_base = row * hidden_size + + # PyTorch's BF16 RMSNorm kernel uses four contiguous values per thread for + # hidden=256. Keep the four ordered BF16 residual values in registers so + # the output pass does not read the large activations a second time. + column_0 = threads * 4 + column_1 = column_0 + 1 + column_2 = column_0 + 2 + column_3 = column_0 + 3 + valid_0 = column_0 < hidden_size + valid_1 = column_1 < hidden_size + valid_2 = column_2 < hidden_size + valid_3 = column_3 < hidden_size + offset_0 = row_base + column_0 + offset_1 = row_base + column_1 + offset_2 = row_base + column_2 + offset_3 = row_base + column_3 + value_0 = _load_ordered_residual(x_ptr, residual_a_ptr, residual_b_ptr, offset_0, valid_0, residual_count) + value_1 = _load_ordered_residual(x_ptr, residual_a_ptr, residual_b_ptr, offset_1, valid_1, residual_count) + value_2 = _load_ordered_residual(x_ptr, residual_a_ptr, residual_b_ptr, offset_2, valid_2, residual_count) + value_3 = _load_ordered_residual(x_ptr, residual_a_ptr, residual_b_ptr, offset_3, valid_3, residual_count) + accumulator = fma_rn_f32(value_0, value_0, tl.zeros((64,), dtype=tl.float32)) + accumulator = fma_rn_f32(value_1, value_1, accumulator) + accumulator = fma_rn_f32(value_2, value_2, accumulator) + accumulator = fma_rn_f32(value_3, value_3, accumulator) + + total = _two_warp_sum(accumulator) + reciprocal_rms = rsqrt_approx_f32(total / hidden_size + eps) + modulation_base = batch * hidden_size + _store_modulated( + output_ptr, + norm_weight_ptr, + scale_ptr, + shift_ptr, + value_0, + offset_0, + column_0, + modulation_base, + reciprocal_rms, + valid_0, + ) + _store_modulated( + output_ptr, + norm_weight_ptr, + scale_ptr, + shift_ptr, + value_1, + offset_1, + column_1, + modulation_base, + reciprocal_rms, + valid_1, + ) + _store_modulated( + output_ptr, + norm_weight_ptr, + scale_ptr, + shift_ptr, + value_2, + offset_2, + column_2, + modulation_base, + reciprocal_rms, + valid_2, + ) + _store_modulated( + output_ptr, + norm_weight_ptr, + scale_ptr, + shift_ptr, + value_3, + offset_3, + column_3, + modulation_base, + reciprocal_rms, + valid_3, + ) + + +@triton.jit +def _residual_add3_kernel( + output_ptr, + x_ptr, + residual_a_ptr, + residual_b_ptr, + residual_c_ptr, + elements, + block: tl.constexpr, +): + offsets = tl.program_id(0).to(tl.int64) * block + tl.arange(0, block) + valid = offsets < elements + value = add_rn_f32( + tl.load(x_ptr + offsets, mask=valid, other=0.0).to(tl.float32), + tl.load(residual_a_ptr + offsets, mask=valid, other=0.0).to(tl.float32), + ) + value = round_bf16_to_fp32(value) + value = add_rn_f32( + value, + tl.load(residual_b_ptr + offsets, mask=valid, other=0.0).to(tl.float32), + ) + value = round_bf16_to_fp32(value) + value = add_rn_f32( + value, + tl.load(residual_c_ptr + offsets, mask=valid, other=0.0).to(tl.float32), + ) + tl.store(output_ptr + offsets, value, mask=valid) + + +def _same_activation(x: torch.Tensor, residual: torch.Tensor) -> bool: + return ( + residual.dtype is torch.bfloat16 + and residual.is_cuda + and residual.device == x.device + and residual.shape == x.shape + and residual.is_contiguous() + ) + + +def _modulation_matches(x: torch.Tensor, value: torch.Tensor) -> bool: + return ( + value.dtype is torch.bfloat16 + and value.is_cuda + and value.device == x.device + and value.shape[-1] == _HIDDEN_SIZE + and value.numel() == x.shape[0] * _HIDDEN_SIZE + and value.is_contiguous() + ) + + +def _norm_weight_matches(x: torch.Tensor, weight: torch.Tensor) -> bool: + return ( + weight.dtype is torch.bfloat16 + and weight.is_cuda + and weight.device == x.device + and weight.shape == (_HIDDEN_SIZE,) + and weight.is_contiguous() + ) + + +def _adaln_inputs_supported( + x: torch.Tensor, + residual_a: torch.Tensor, + residual_b: torch.Tensor | None, + norm_weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + eps: float, +) -> bool: + return ( + is_ltx2_fusion_eligible(x) + and x.dtype is torch.bfloat16 + and x.ndim >= 2 + and x.shape[-1] == _HIDDEN_SIZE + and x.numel() > 0 + and x.is_contiguous() + and _same_activation(x, residual_a) + and (residual_b is None or _same_activation(x, residual_b)) + and _norm_weight_matches(x, norm_weight) + and _modulation_matches(x, scale) + and _modulation_matches(x, shift) + and math.isfinite(eps) + and eps > 0 + ) + + +def _residual_reference( + x: torch.Tensor, + residual_a: torch.Tensor, + residual_b: torch.Tensor | None, +) -> torch.Tensor: + hidden_states = x + residual_a + if residual_b is not None: + hidden_states = hidden_states + residual_b + return hidden_states + + +def _adaln_reference( + x: torch.Tensor, + residual_a: torch.Tensor, + residual_b: torch.Tensor | None, + norm_weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + eps: float, +) -> torch.Tensor: + hidden_states = _residual_reference(x, residual_a, residual_b) + normalized = F.rms_norm(hidden_states, (_HIDDEN_SIZE,), weight=norm_weight, eps=eps) + return normalized * (1 + scale) + shift + + +def _launch_adaln( + x: torch.Tensor, + residual_a: torch.Tensor, + residual_b: torch.Tensor | None, + norm_weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + eps: float, +) -> torch.Tensor: + output = torch.empty_like(x) + rows = x.numel() // _HIDDEN_SIZE + rows_per_batch = rows // x.shape[0] + with torch.accelerator.device_index(x.device.index): + _residual_rms_norm_modulate_kernel[(rows,)]( + output, + x, + residual_a, + x if residual_b is None else residual_b, + norm_weight, + scale, + shift, + rows_per_batch, + eps, + residual_count=1 if residual_b is None else 2, + hidden_size=_HIDDEN_SIZE, + num_warps=2, + ) + return output + + +def _verify_adaln_prefix( + output: torch.Tensor, + x: torch.Tensor, + residual_a: torch.Tensor, + residual_b: torch.Tensor | None, + norm_weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + eps: float, +) -> bool: + rows_per_batch = x.numel() // (x.shape[0] * _HIDDEN_SIZE) + rows = min(rows_per_batch, _VERIFY_ROWS) + x_rows = x.reshape(x.shape[0], rows_per_batch, _HIDDEN_SIZE)[0, :rows] + residual_a_rows = residual_a.reshape(x.shape[0], rows_per_batch, _HIDDEN_SIZE)[0, :rows] + residual_b_rows = ( + None if residual_b is None else residual_b.reshape(x.shape[0], rows_per_batch, _HIDDEN_SIZE)[0, :rows] + ) + scale_row = scale.reshape(x.shape[0], 1, _HIDDEN_SIZE)[0] + shift_row = shift.reshape(x.shape[0], 1, _HIDDEN_SIZE)[0] + reference = _adaln_reference( + x_rows, + residual_a_rows, + residual_b_rows, + norm_weight, + scale_row, + shift_row, + eps, + ) + actual = output.reshape(x.shape[0], rows_per_batch, _HIDDEN_SIZE)[0, :rows] + return torch.equal(actual, reference) + + +def try_residual_rms_norm_modulate_exact( + x: torch.Tensor, + residual_a: torch.Tensor, + residual_b: torch.Tensor | None, + norm_weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + eps: float, +) -> torch.Tensor | None: + """Fuse one or two ordered BF16 residuals with RMSNorm and AdaLN.""" + + if not _adaln_inputs_supported(x, residual_a, residual_b, norm_weight, scale, shift, eps): + return None + residual_count = 1 if residual_b is None else 2 + runtime_key = (x.device.index, residual_count, float(eps)) + if runtime_key in _FAILED_ADALN_KEYS: + return None + try: + output = _launch_adaln(x, residual_a, residual_b, norm_weight, scale, shift, eps) + if runtime_key not in _VERIFIED_ADALN_KEYS: + if not _verify_adaln_prefix( + output, + x, + residual_a, + residual_b, + norm_weight, + scale, + shift, + eps, + ): + _FAILED_ADALN_KEYS.add(runtime_key) + logger.warning( + "Disabling LTX DiffVAE residual-AdaLN fusion on %s after a bit-exactness mismatch", + x.device, + ) + return None + _VERIFIED_ADALN_KEYS.add(runtime_key) + except Exception as exc: # noqa: BLE001 - fail closed after optimized-path failure + _FAILED_ADALN_KEYS.add(runtime_key) + logger.warning( + "Disabling LTX DiffVAE residual-AdaLN fusion on %s after failure: %s", + x.device, + exc, + ) + return None + return output + + +def _add_inputs_supported(x: torch.Tensor, residuals: tuple[torch.Tensor, ...]) -> bool: + return ( + is_ltx2_fusion_eligible(x) + and x.dtype is torch.bfloat16 + and x.numel() > 0 + and x.is_contiguous() + and all(_same_activation(x, residual) for residual in residuals) + ) + + +def _launch_add3( + x: torch.Tensor, + residual_a: torch.Tensor, + residual_b: torch.Tensor, + residual_c: torch.Tensor, +) -> torch.Tensor: + output = torch.empty_like(x) + with torch.accelerator.device_index(x.device.index): + _residual_add3_kernel[(triton.cdiv(x.numel(), _POINTWISE_BLOCK),)]( + output, + x, + residual_a, + residual_b, + residual_c, + x.numel(), + block=_POINTWISE_BLOCK, + num_warps=4, + ) + return output + + +def try_residual_add3_exact( + x: torch.Tensor, + residual_a: torch.Tensor, + residual_b: torch.Tensor, + residual_c: torch.Tensor, +) -> torch.Tensor | None: + """Fuse three ordered BF16 residual additions into one pointwise pass.""" + + residuals = (residual_a, residual_b, residual_c) + if not _add_inputs_supported(x, residuals) or x.device.index in _FAILED_ADD_DEVICES: + return None + try: + output = _launch_add3(x, *residuals) + if x.device.index not in _VERIFIED_ADD_DEVICES: + elements = min(x.numel(), _VERIFY_ELEMENTS) + reference = ( + _residual_reference( + x.reshape(-1)[:elements], + residual_a.reshape(-1)[:elements], + residual_b.reshape(-1)[:elements], + ) + + residual_c.reshape(-1)[:elements] + ) + if not torch.equal(output.reshape(-1)[:elements], reference): + _FAILED_ADD_DEVICES.add(x.device.index) + logger.warning( + "Disabling LTX DiffVAE residual-add fusion on %s after a bit-exactness mismatch", + x.device, + ) + return None + _VERIFIED_ADD_DEVICES.add(x.device.index) + except Exception as exc: # noqa: BLE001 - fail closed after optimized-path failure + _FAILED_ADD_DEVICES.add(x.device.index) + logger.warning( + "Disabling LTX DiffVAE residual-add fusion on %s after failure: %s", + x.device, + exc, + ) + return None + return output + + +__all__ = [ + "try_residual_add3_exact", + "try_residual_rms_norm_modulate_exact", +] diff --git a/vllm_omni/diffusion/models/ltx2/ops/rope_tables.py b/vllm_omni/diffusion/models/ltx2/ops/rope_tables.py new file mode 100644 index 000000000..78a2ad2ee --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ops/rope_tables.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Cached 3D RoPE tables for LTX DiffVAE kernels.""" + +from __future__ import annotations + +import torch + +_TABLE_CACHE_MAX = 16 +_TABLE_CACHE: dict[ + tuple[torch.device, int, int, int, tuple[int, int, int], float], + tuple[tuple[torch.Tensor, torch.Tensor], ...], +] = {} + + +def _axis_tables( + length: int, + dim: int, + device: torch.device, + base: float, +) -> tuple[torch.Tensor, torch.Tensor]: + exponents = torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim + inv_freqs = (1.0 / base**exponents).to(torch.float32) + positions = torch.arange(length, dtype=torch.float32, device=device) + angles = positions[:, None] * inv_freqs[None, :] + return angles.cos().contiguous(), angles.sin().contiguous() + + +def get_rope_tables( + query: torch.Tensor, + dim_split: tuple[int, int, int], + base: float, +) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: + num_frames, height, width = query.shape[1:4] + cache_key = (query.device, num_frames, height, width, dim_split, base) + cached = _TABLE_CACHE.get(cache_key) + if cached is not None: + return cached + tables = tuple( + _axis_tables(length, dim, query.device, base) + for length, dim in zip((num_frames, height, width), dim_split, strict=True) + ) + if len(_TABLE_CACHE) >= _TABLE_CACHE_MAX: + _TABLE_CACHE.pop(next(iter(_TABLE_CACHE))) + _TABLE_CACHE[cache_key] = tables + return tables + + +__all__ = ["get_rope_tables"] diff --git a/vllm_omni/diffusion/models/ltx2/ops/swiglu.py b/vllm_omni/diffusion/models/ltx2/ops/swiglu.py new file mode 100644 index 000000000..6c332b7ed --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ops/swiglu.py @@ -0,0 +1,206 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Bit-exact workspace-reusing SwiGLU for the LTX diffusion VAE decoder.""" + +from __future__ import annotations + +import logging + +import torch +import torch.nn.functional as F +from vllm.triton_utils import tl, triton + +from . import is_ltx2_fusion_eligible +from ._numerics import mul_rn_f32, round_bf16_to_fp32 + +_POINTWISE_BLOCK = 1024 +_BASE_TILE_ROWS = 16384 +_MAX_TILE_HIDDEN_ELEMENTS = _BASE_TILE_ROWS * 4 * 2048 +_FAILED_KEYS: set[tuple[int | None, int, int, int, int]] = set() +_VERIFIED_KEYS: set[tuple[int | None, int, int, int, int]] = set() + +logger = logging.getLogger(__name__) + + +def _tile_rows(rows: int, hidden: int) -> int: + """Bound projection storage while giving narrow MLPs larger GEMMs.""" + + budget_rows = _MAX_TILE_HIDDEN_ELEMENTS // hidden + aligned_rows = max(_BASE_TILE_ROWS, budget_rows // _BASE_TILE_ROWS * _BASE_TILE_ROWS) + return min(rows, aligned_rows) + + +@triton.jit +def _silu_mul_kernel( + gate_ptr, + up_ptr, + elements, + block: tl.constexpr, +): + offsets = tl.program_id(0).to(tl.int64) * block + tl.arange(0, block) + valid = offsets < elements + gate = tl.load(gate_ptr + offsets, mask=valid).to(tl.float32) + up = tl.load(up_ptr + offsets, mask=valid).to(tl.float32) + # Eager rounds SiLU to BF16 before the separate BF16 multiply. + silu = round_bf16_to_fp32(mul_rn_f32(gate, tl.sigmoid(gate))) + tl.store(gate_ptr + offsets, mul_rn_f32(silu, up), mask=valid) + + +def _weight_matches( + hidden_states: torch.Tensor, + weight: torch.Tensor, + shape: tuple[int, int], +) -> bool: + return ( + weight.is_cuda + and weight.device == hidden_states.device + and weight.dtype is torch.bfloat16 + and weight.shape == shape + and weight.is_contiguous() + ) + + +def _supported_inputs( + hidden_states: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, +) -> bool: + if hidden_states.ndim < 3: + return False + dim = hidden_states.shape[-1] + hidden = gate_weight.shape[0] if gate_weight.ndim == 2 else 0 + return ( + is_ltx2_fusion_eligible(hidden_states) + and hidden_states.dtype is torch.bfloat16 + and hidden_states.shape[0] == 1 + and hidden_states.numel() > 0 + and hidden_states.is_contiguous() + and 256 <= dim <= 2048 + and hidden == 4 * dim + and _weight_matches(hidden_states, gate_weight, (hidden, dim)) + and _weight_matches(hidden_states, up_weight, (hidden, dim)) + and _weight_matches(hidden_states, down_weight, (dim, hidden)) + ) + + +def _silu_mul_inplace(gate: torch.Tensor, up: torch.Tensor) -> None: + with torch.accelerator.device_index(gate.device.index): + _silu_mul_kernel[(triton.cdiv(gate.numel(), _POINTWISE_BLOCK),)]( + gate, + up, + gate.numel(), + block=_POINTWISE_BLOCK, + num_warps=4, + ) + + +def _reference_tile( + hidden_states: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, +) -> torch.Tensor: + return F.linear( + F.silu(F.linear(hidden_states, gate_weight)) * F.linear(hidden_states, up_weight), + down_weight, + ) + + +def _launch( + hidden_states: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, +) -> torch.Tensor: + dim = hidden_states.shape[-1] + hidden = gate_weight.shape[0] + flat = hidden_states.reshape(-1, dim) + tile_rows = _tile_rows(flat.shape[0], hidden) + output = torch.empty_like(flat) + workspace_gate = torch.empty((tile_rows, hidden), device=flat.device, dtype=flat.dtype) + for start in range(0, flat.shape[0], tile_rows): + end = min(start + tile_rows, flat.shape[0]) + tokens = end - start + tile = flat[start:end] + gate = workspace_gate[:tokens] + torch.mm(tile, gate_weight.t(), out=gate) + up = torch.mm(tile, up_weight.t()) + _silu_mul_inplace(gate, up) + torch.mm(gate, down_weight.t(), out=output[start:end]) + return output.reshape(hidden_states.shape) + + +def _verify_edge_tiles( + output: torch.Tensor, + hidden_states: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, +) -> bool: + from diffusers.models.autoencoders.ltx2_diffusion_decoder import _SWIGLU_TILE_SIZE + + dim = hidden_states.shape[-1] + flat = hidden_states.reshape(-1, dim) + actual = output.reshape(-1, dim) + tile_rows = _tile_rows(flat.shape[0], gate_weight.shape[0]) + starts = {0, ((flat.shape[0] - 1) // tile_rows) * tile_rows} + for start in starts: + end = min(start + tile_rows, flat.shape[0]) + # Match upstream's token tiling, not the optimized GEMM shape: changing + # the matmul row count can change its reduction order and BF16 output. + for reference_start in range(start, end, _SWIGLU_TILE_SIZE): + reference_end = min(reference_start + _SWIGLU_TILE_SIZE, end) + if not torch.equal( + actual[reference_start:reference_end], + _reference_tile(flat[reference_start:reference_end], gate_weight, up_weight, down_weight), + ): + return False + return True + + +def try_swiglu_tiled_exact( + hidden_states: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, +) -> torch.Tensor | None: + """Run workspace-reusing tiled SwiGLU, or expose the original module.""" + + if not _supported_inputs(hidden_states, gate_weight, up_weight, down_weight): + return None + rows = hidden_states.numel() // hidden_states.shape[-1] + tile_rows = _tile_rows(rows, gate_weight.shape[0]) + runtime_key = ( + hidden_states.device.index, + hidden_states.shape[-1], + gate_weight.shape[0], + tile_rows, + (rows - 1) % tile_rows + 1, + ) + if runtime_key in _FAILED_KEYS: + return None + try: + output = _launch(hidden_states, gate_weight, up_weight, down_weight) + if runtime_key not in _VERIFIED_KEYS: + if not _verify_edge_tiles(output, hidden_states, gate_weight, up_weight, down_weight): + _FAILED_KEYS.add(runtime_key) + logger.warning( + "Disabling LTX DiffVAE tiled SwiGLU fusion on %s after a bit-exactness mismatch", + hidden_states.device, + ) + return None + _VERIFIED_KEYS.add(runtime_key) + except Exception as exc: # noqa: BLE001 - fail closed after optimized-path failure + _FAILED_KEYS.add(runtime_key) + logger.warning( + "Disabling LTX DiffVAE tiled SwiGLU fusion on %s after failure: %s", + hidden_states.device, + exc, + ) + return None + return output + + +__all__ = ["try_swiglu_tiled_exact"] diff --git a/vllm_omni/diffusion/models/ltx2/vae/__init__.py b/vllm_omni/diffusion/models/ltx2/vae/__init__.py new file mode 100644 index 000000000..9381766fc --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/vae/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Ordinary LTX-2.5 DiffVAE, tiled decoding and acceleration operators.""" diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder.py b/vllm_omni/diffusion/models/ltx2/vae/decoder.py similarity index 57% rename from vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder.py rename to vllm_omni/diffusion/models/ltx2/vae/decoder.py index 73db2ea6b..312578990 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder.py +++ b/vllm_omni/diffusion/models/ltx2/vae/decoder.py @@ -7,19 +7,24 @@ """LTX-2.5 Native-weight adapter for Diffusers' diffusion VAE decoder.""" +import math from collections.abc import Mapping +from functools import cache from typing import Any import torch from diffusers.models.autoencoders import ( LTX2VideoDiffusionDecoderModel as DiffusersLTX2VideoDiffusionDecoderModel, ) +from diffusers.models.autoencoders import ltx2_diffusion_decoder as diffusers_decoder from diffusers.models.autoencoders.ltx2_diffusion_decoder import ( LTX2VideoDiffusionDecoder3d as DiffusersLTX2VideoDiffusionDecoder3d, ) from diffusers.models.autoencoders.vae import DecoderOutput from safetensors import safe_open +from ..ops import is_ltx2_fna_eligible, is_ltx2_fusion_eligible, resolve_ltx2_vae_operators + LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID = "Lightricks/LTX-2.5" LTX25_NATIVE_ARTIFACT_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1" LTX25_NATIVE_DIFFUSION_DECODER_FILENAME = "vae/ltx-2.5-video-vae-bf16.safetensors" @@ -131,6 +136,177 @@ def load_ltx25_native_diffusion_decoder_state_dict(path: str) -> dict[str, torch return convert_ltx25_native_diffusion_decoder_state_dict(native_state_dict) +@cache +def _load_natten_token_permutation(): + from kernels import get_kernel + + natten = get_kernel("shi-labs/natten", version=1) + return natten.token_permute.token_permute_operation, natten.token_permute.token_unpermute_operation + + +class LTX2VideoVaeNeighborhoodNattenProcessor(diffusers_decoder.LTX2VideoVaeNeighborhoodNattenProcessor): + """NATTEN dispatch with Omni FNA for eligible stage-5 grids.""" + + def __call__( + self, attn: "LTX2VideoVaeNeighborhoodAttention", hidden_states: torch.Tensor, block_mask=None + ) -> torch.Tensor: + batch_size, num_frames, height, width, channels = hidden_states.shape + query, key, value = attn.project_qkv(hidden_states) + # NATTEN's CUTLASS kernels silently produce wrong output for non-contiguous inputs. + query, key, value = query.contiguous(), key.contiguous(), value.contiguous() + operators = resolve_ltx2_vae_operators(hidden_states.device) if is_ltx2_fna_eligible(hidden_states) else None + # `scale=1.0`: the query is already scaled in `project_qkv`, as in the reference. + if ( + tuple(attn.kernel_size) == (11, 11, 11) + and hidden_states.dtype == torch.bfloat16 + and channels == 256 + and attn.head_dim == 64 + and operators is not None + and operators.fna is not None + # The native metadata encodes KV tile indices in 15 bits. + and math.prod((n + tile - 1) // tile for n, tile in zip((num_frames, height, width), (4, 4, 8))) <= 1 << 15 + ): + token_permute, token_unpermute = _load_natten_token_permutation() + query_permuted, token_shape, _ = token_permute(query, (4, 4, 4), flip_tiled_dims=True) + key_permuted, _, _ = token_permute(key, (4, 4, 8), flip_tiled_dims=True) + value_permuted, _, _ = token_permute(value, (4, 4, 8), flip_tiled_dims=True) + output = operators.fna( + query_permuted, + key_permuted, + value_permuted, + shape=(num_frames, height, width), + window=tuple(attn.kernel_size), + ) + hidden_states = token_unpermute( + output, + token_layout_shape=token_shape, + tile_shape=(4, 4, 4), + flip_tiled_dims=True, + ) + else: + hidden_states = self._na3d( + query, + key, + value, + kernel_size=attn.kernel_size, + scale=1.0, + backend=self.backend, + ) + hidden_states = hidden_states.reshape(batch_size, num_frames, height, width, channels) + return attn.to_out[0](hidden_states) + + +class LTX2VideoVaeNeighborhoodAttention(diffusers_decoder.LTX2VideoVaeNeighborhoodAttention): + _available_processors = [ + *diffusers_decoder.LTX2VideoVaeNeighborhoodAttention._available_processors, + LTX2VideoVaeNeighborhoodNattenProcessor, + ] + + def project_qkv(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size, num_frames, height, width, _ = hidden_states.shape + shape = (batch_size, num_frames, height, width, self.heads, self.head_dim) + query = self.to_q(hidden_states).view(shape) + key = self.to_k(hidden_states).view(shape) + value = self.to_v(hidden_states).view(shape) + + operators = resolve_ltx2_vae_operators(query.device) if is_ltx2_fusion_eligible(query) else None + optimized = ( + None + if operators is None + else operators.qk_norm_rope( + query, + key, + self.norm_q.weight, + self.norm_k.weight, + self.norm_q.eps, + self.scale, + self.rope.rope_dim_split, + self.rope.base, + ) + ) + if optimized is not None: + query, key = optimized + return query, key, value + + query = self.norm_q(query) + key = self.norm_k(key) + query = query * self.scale + return self.rope(query), self.rope(key), value + + +class LTX2VideoVaeSwiGLU(diffusers_decoder.LTX2VideoVaeSwiGLU): + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + operators = resolve_ltx2_vae_operators(hidden_states.device) if is_ltx2_fusion_eligible(hidden_states) else None + optimized = ( + None + if operators is None + else operators.swiglu(hidden_states, self.w_gate.weight, self.w_up.weight, self.w_down.weight) + ) + return super().forward(hidden_states) if optimized is None else optimized + + +class LTX2VideoVaeDiffusionNABlock(diffusers_decoder.LTX2VideoVaeDiffusionNABlock): + def forward( + self, + hidden_states: torch.Tensor, + latent_context: torch.Tensor, + modulation: tuple[torch.Tensor, ...], + block_mask=None, + ) -> torch.Tensor: + if hidden_states.dtype is not torch.bfloat16 or not is_ltx2_fusion_eligible(hidden_states): + return super().forward(hidden_states, latent_context, modulation, block_mask) + operators = resolve_ltx2_vae_operators(hidden_states.device) + assert operators is not None + + scale_msa, shift_msa, _, scale_mlp, shift_mlp, _, _ = [ + modulation[i] + self.scale_shift_table[i].view(1, 1, 1, 1, -1) for i in range(self.num_mod_params) + ] + context_output = self.context_proj(latent_context) + attention_input = operators.residual_norm( + hidden_states, + context_output, + None, + self.norm1.weight, + scale_msa, + shift_msa, + self.norm1.eps, + ) + if attention_input is None: + hidden_states = hidden_states + context_output + hidden_states = hidden_states + self.attn( + self.norm1(hidden_states) * (1 + scale_msa) + shift_msa, block_mask + ) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states) * (1 + scale_mlp) + shift_mlp) + return hidden_states + + attention_output = self.attn(attention_input, block_mask) + mlp_input = operators.residual_norm( + hidden_states, + context_output, + attention_output, + self.norm2.weight, + scale_mlp, + shift_mlp, + self.norm2.eps, + ) + if mlp_input is None: + residual = hidden_states + context_output + residual = residual + attention_output + mlp_input = self.norm2(residual) * (1 + scale_mlp) + shift_mlp + mlp_output = self.mlp(mlp_input) + output = operators.residual_add( + hidden_states, + context_output, + attention_output, + mlp_output, + ) + if output is not None: + return output + residual = hidden_states + context_output + residual = residual + attention_output + return residual + mlp_output + + class LTX2VideoDiffusionDecoder3d(DiffusersLTX2VideoDiffusionDecoder3d): """Diffusers decoder core with the short-clip NATTEN context fix.""" @@ -200,6 +376,15 @@ class LTX2VideoDiffusionDecoderModel(DiffusersLTX2VideoDiffusionDecoderModel): # to this behavior-only subclass preserves parameters and state-dict keys. self.decoder.__class__ = LTX2VideoDiffusionDecoder3d self.decoder.stage5_kernel = tuple(decoder_stage5_kernel) + overrides = { + diffusers_decoder.LTX2VideoVaeNeighborhoodAttention: LTX2VideoVaeNeighborhoodAttention, + diffusers_decoder.LTX2VideoVaeSwiGLU: LTX2VideoVaeSwiGLU, + diffusers_decoder.LTX2VideoVaeDiffusionNABlock: LTX2VideoVaeDiffusionNABlock, + } + # Preserve every module and parameter object; only execution methods change. + for module in self.decoder.modules(): + if replacement := overrides.get(type(module)): + module.__class__ = replacement @classmethod def from_ltx25_native_checkpoint( diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder_distributed.py b/vllm_omni/diffusion/models/ltx2/vae/distributed.py similarity index 99% rename from vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder_distributed.py rename to vllm_omni/diffusion/models/ltx2/vae/distributed.py index 7f486dbdf..9a98ba739 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder_distributed.py +++ b/vllm_omni/diffusion/models/ltx2/vae/distributed.py @@ -25,7 +25,7 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor TileTask, ) -from .ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel +from .decoder import LTX2VideoDiffusionDecoderModel logger = logging.get_logger(__name__) # pylint: disable=invalid-name