410 lines
15 KiB
Python
410 lines
15 KiB
Python
# 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
|