90 lines
3.4 KiB
Python
90 lines
3.4 KiB
Python
# 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))
|