# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import math import sys import types from typing import Any import pytest import torch from vllm_omni.diffusion.attention.backends.abstract import ( AttentionMetadata, PackedPaddingMetadata, VideoTokenLayout, VideoTokenSpan, ) from vllm_omni.diffusion.models.minimax_h3.attention.fastvideo_h3 import ( MiniMaxH3VSAImpl, _get_h3_tile_metadata, ) pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None): extra: dict[str, Any] = { "vsa_h3_prefix_segments": prefix_segments, } if gate is not None: extra["gate_compress"] = gate return AttentionMetadata( packed_padding=packed_padding, video_layout=VideoTokenLayout( used_len=sum(prefix_segments) + math.prod(video_shape), video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),), ), extra=extra, ) def _fake_block_sparse_kernel(monkeypatch, calls): """Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA.""" def block_sparse_attn(q, k, v, block_map, variable_block_sizes): calls["block_map"] = block_map.clone() calls["variable_block_sizes"] = variable_block_sizes.tolist() calls["q_shape"] = tuple(q.shape) blocks = block_map.shape[-1] assert blocks * 64 == q.shape[2] # Expand the per-block map to a per-token mask so the stand-in honours # the sparsity contract the real kernel implements, and drop the rows # each partial block pads with, which the real kernel skips through # variable_block_sizes. token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2) within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None] token_mask = token_mask & within.reshape(1, 1, 1, -1) out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask) return out.to(q.dtype), None module = types.ModuleType("fastvideo_kernel.block_sparse_attn") setattr(module, "block_sparse_attn", block_sparse_attn) monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel")) monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module) def _h3_impl(**backend_kwargs): kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False} kwargs.update(backend_kwargs) return MiniMaxH3VSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs) def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch): calls: dict[str, Any] = {} _fake_block_sparse_kernel(monkeypatch, calls) impl = _h3_impl(topk=64) prefix_segments, video_shape = (5, 70), (2, 4, 4) seq_len = sum(prefix_segments) + 2 * 4 * 4 torch.manual_seed(0) query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16) output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) assert output.shape == query.shape # topk covers every video block, so the block map is dense and the tiled # kernel must reproduce plain attention over the same rows. reference = torch.nn.functional.scaled_dot_product_attention( *(x.float().transpose(1, 2) for x in (query, query, query)) ).transpose(1, 2) torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2) # 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile. assert calls["variable_block_sizes"] == [5, 64, 6, 32] def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch): calls: dict[str, Any] = {} _fake_block_sparse_kernel(monkeypatch, calls) impl = _h3_impl(topk=1) prefix_segments, video_shape = (5,), (8, 4, 4) seq_len = 5 + 8 * 4 * 4 torch.manual_seed(0) query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16) impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) # 1 prefix block + 2 video blocks. The odd count is padded to an even one # for the paired-CTA kernel contract, and the Triton route slices that # transport-only partner off again before launching. block_map = calls["block_map"] assert block_map.shape[-1] == 3 assert calls["variable_block_sizes"] == [5, 64, 64] assert block_map[:, :, :1, :].all(), "prefix queries stay dense" assert block_map[..., :1].all(), "prefix keys are exempt from selection" assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query" # The same odd geometry must still untile back onto the original rows. dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) reference = torch.nn.functional.scaled_dot_product_attention( *(x.float().transpose(1, 2) for x in (query, query, query)) ).transpose(1, 2) torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2) def test_h3_forward_applies_the_learned_compression_gate(monkeypatch): _fake_block_sparse_kernel(monkeypatch, {}) impl = _h3_impl(topk=64) prefix_segments, video_shape = (5,), (2, 4, 4) seq_len = 5 + 2 * 4 * 4 torch.manual_seed(0) query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16) without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) zero_gate = torch.zeros_like(query) with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate)) live_gate = torch.full_like(query, 0.5) with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate)) torch.testing.assert_close(with_zero_gate.float(), without_gate.float()) assert not torch.allclose(with_live_gate.float(), without_gate.float()) def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch): _fake_block_sparse_kernel(monkeypatch, {}) impl = _h3_impl(topk=64) prefix_segments, video_shape = (5,), (2, 4, 4) valid = 5 + 2 * 4 * 4 torch.manual_seed(0) query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16) padding = PackedPaddingMetadata( q_length=valid, kv_length=valid, cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32), cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32), ) output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding)) assert output.shape == query.shape assert torch.count_nonzero(output[:, valid:]) == 0 def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d(): partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata( (5, 70, 9), (5, 6, 6), torch.device("cpu") ) assert sizes[:4].tolist() == [5, 64, 6, 9] assert prefix_blocks == 4 assert video_blocks == 8 assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6 assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum()) @pytest.mark.parametrize("fallback_on_error", [False, True]) def test_h3_kernel_failure_preserves_fallback_policy(monkeypatch, fallback_on_error): impl = _h3_impl(fallback_on_error=fallback_on_error) query = torch.randn(1, 37, 2, 8) metadata = _h3_metadata((5,), (2, 4, 4)) def fail(*args, **kwargs): raise RuntimeError("kernel failure") monkeypatch.setattr(impl, "_forward_h3", fail) monkeypatch.setattr(impl.sdpa_fallback, "forward", impl.sdpa_fallback.forward_cuda) if fallback_on_error: output = impl.forward_cuda(query, query, query, metadata) expected = impl.sdpa_fallback.forward_cuda(query, query, query, metadata) torch.testing.assert_close(output, expected, rtol=0, atol=0) else: with pytest.raises(RuntimeError, match="kernel failure"): impl.forward_cuda(query, query, query, metadata) def test_h3_accelerator_failure_is_not_swallowed(monkeypatch): impl = _h3_impl(fallback_on_error=True) query = torch.randn(1, 37, 2, 8) def fail(*args, **kwargs): raise torch.AcceleratorError("device failure") monkeypatch.setattr(impl, "_forward_h3", fail) with pytest.raises(torch.AcceleratorError, match="device failure"): impl.forward_cuda(query, query, query, _h3_metadata((5,), (2, 4, 4))) def test_h3_without_layout_delegates_to_shared_backend(monkeypatch): from vllm_omni.diffusion.attention.backends.fastvideo_vsa import FastVideoVSAImpl impl = _h3_impl() query = torch.randn(1, 37, 2, 8) metadata = AttentionMetadata() expected = torch.empty_like(query) calls = [] def forward(self, q, k, v, attn_metadata): calls.append((self, q, k, v, attn_metadata)) return expected monkeypatch.setattr(FastVideoVSAImpl, "forward_cuda", forward) assert impl.forward_cuda(query, query, query, metadata) is expected assert len(calls) == 1 assert all(actual is wanted for actual, wanted in zip(calls[0], (impl, query, query, query, metadata)))